Skip to content

Instantly share code, notes, and snippets.

@vadimkantorov

vadimkantorov/compact_bilinear_pooling.py

Last active September 22, 2021 07:51
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save vadimkantorov/d9b56f9b85f1f4ce59ffecf893a1581a to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save vadimkantorov/d9b56f9b85f1f4ce59ffecf893a1581a to your computer and use it in GitHub Desktop.
Download ZIP
Compact Bilinear Pooling in PyTorch using the new FFT support
Raw
compact_bilinear_pooling.py
import torch
import torch.nn as nn
class CompactBilinearPooling(nn.Module):
def __init__(self, input_dim1, input_dim2, output_dim, sum_pool = True):
super(CompactBilinearPooling, self).__init__()
self.output_dim = output_dim
self.sum_pool = sum_pool
generate_sketch_matrix = lambda rand_h, rand_s, input_dim, output_dim: torch.sparse.FloatTensor(torch.stack([torch.arange(input_dim, out = torch.LongTensor()), rand_h.long()]), rand_s.float(), torch.Size([input_dim, output_dim])).to_dense()
self.sketch_matrix1 = nn.Parameter(generate_sketch_matrix(torch.randint(output_dim, size = (input_dim1,)), 2 * torch.randint(2, size = (input_dim1,)) - 1, input_dim1, output_dim))
self.sketch_matrix2 = nn.Parameter(generate_sketch_matrix(torch.randint(output_dim, size = (input_dim2,)), 2 * torch.randint(2, size = (input_dim2,)) - 1, input_dim2, output_dim))
def forward(self, bottom1, bottom2):
sketch_1 = bottom1.permute(0, 2, 3, 1).contiguous().matmul(self.sketch_matrix1).view(-1, self.output_dim)
sketch_2 = bottom2.permute(0, 2, 3, 1).contiguous().matmul(self.sketch_matrix2).view(-1, self.output_dim)
fft1_real, fft1_imag = torch.rfft(sketch_1, 1).permute(2, 0, 1)
fft2_real, fft2_imag = torch.rfft(sketch_2, 1).permute(2, 0, 1)
fft_product = torch.stack([fft1_real * fft2_real - fft1_imag * fft2_imag, fft1_real * fft2_imag - fft1_imag * fft2_real], dim = -1)
cbp = torch.irfft(fft_product, 1, signal_sizes = (self.output_dim,)).view(len(bottom1), bottom1.size(-2), bottom1.size(-1), self.output_dim) * self.output_dim
return cbp.sum(dim = 1).sum(dim = 1) if self.sum_pool else cbp.permute(0, 3, 1, 2)
@pangjh3
Copy link

pangjh3 commented Apr 20, 2018

Thanks for your code, how to install the new fft support?

@vadimkantorov
Copy link
Author

vadimkantorov commented Jun 5, 2018

Just install PyTorch from master branch or even 0.4 version probably has FFT

@ayumiymk
Copy link

ayumiymk commented Dec 12, 2018

Thanks for your code first. I have a question that in the other implements, like Torch version and Tensorflow version, there is a zero_padding before feeding the tensor into the fft. But in this code, I don't see the zero_padding.

Thanks very much!

@hj0921
Copy link

hj0921 commented Mar 19, 2021

hello,

torch.stack([torch.arange(in_features), rand_h]) where in_features is not defined. How to fix it?

thanks!

@vadimkantorov
Copy link
Author

vadimkantorov commented Mar 19, 2021

Thanks for noting this. Fixed! It should have been in_channels

@vadimkantorov
Copy link
Author

vadimkantorov commented Mar 19, 2021

Some ways to improve the code: make use of the new PyTorch fft module, complex support. Figure out dense x sparse matmul (currently I'm materializing the sparse sketch)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment