def forward(self, k_re, k_im):
k_re, k_im = make_contiguous(k_re, k_im)
x_re, x_im = ifft(k_re, k_im)
if self.norm == 'ortho':
N = np.sqrt(x_re.size(-1))
x_re *= N
x_im *= N
return x_re, x_im
def backward(self, grad_output_re, grad_output_im):
grad_output_re, grad_output_im = make_contiguous(
grad_output_re, grad_output_im)
gi, gr = ifft(grad_output_im, grad_output_re)
if self.norm == 'ortho':
N = np.sqrt(gi.size(-1))
gi *= N
gr *= N
return gr, gi
def conv(self, x, y):
batchSize = x.size(0)
dim = x.size(1)
x_i = torch.FloatTensor(x.size()).cuda().zero_()
y_i = torch.FloatTensor(x.size()).cuda().zero_()
x1_r, x1_i = cfft.fft(x, x_i)
y1_r, y1_i = cfft.fft(y, y_i)
return cfft.ifft(x1_r * y1_r, x1_i * y1_i)
def step(self, closure=None):
"""
Performs a single optimization step.
Arguments:
closure (callable, optional): A closure that reevaluates the model and returns the loss.
"""
loss = None
if closure is not None:
loss = closure()
#import ipdb; ipdb.set_trace()
for group in self.param_groups:
weight_decay = group['weight_decay']
momentum = group['momentum']
dampening = group['dampening']
nesterov = group['nesterov']
# Update the parameters
idx = 0
for param in group['params']:
if param.grad is None:
continue
tmp = param.grad.view(-1, self.sizes[idx])
tmp = tmp.data
re, im = fft.fft(tmp, self.zero_Ns[idx])
re = re*self.coeffs[idx]
im = im*self.coeffs[idx]
tmp = fft.ifft(re, im)[0]
tmp = tmp.view(param.grad.size())
param.grad.data = tmp
idx += 1
d_p = param.grad.data
if weight_decay != 0:
d_p.add_(weight_decay, param.data)
if momentum != 0:
param_state = self.state[param]
if 'momentum_buffer' not in param_state:
buf = param_state['momentum_buffer'] = torch.zeros_like(param.data)
buf.mul_(momentum).add_(d_p)
else:
buf = param_state['momentum_buffer']
buf.mul_(momentum).add_(1 - dampening, d_p)
if nesterov:
d_p = d_p.add(momentum, buf)
else:
d_p = buf
param.data.add_(-group['lr'], d_p)
return loss
def run_fft1(x, z):
if torch.cuda.is_available():
y1, y2 = cfft.fft(x, z)
x_np = x.cpu().numpy().squeeze()
y_np = nfft.fft(x_np)
# print(y1.cpu().numpy())
# print(y_np.real)
assert np.allclose(y1.cpu().numpy(), y_np.real)
assert np.allclose(y2.cpu().numpy(), y_np.imag)
# assert np.allclose(y1[1,0].cpu().numpy(), nfft.fft2(x_np[1,0]).real)
x0, z0 = cfft.ifft(y1, y2)
x0_np = nfft.ifft(y_np)
assert np.allclose(x0.cpu().numpy(), x0_np.real)
assert np.allclose(z0.cpu().numpy(), x0_np.imag)
else:
print("Cuda not available, cannot test.")
def test_acc3():
batch = 3
nch = 4
n = 5
x = torch.randn(batch * nch * n).view(batch, nch, n).cuda()
y = torch.randn(batch * nch * n).view(batch, nch, n).cuda()
x_i = torch.zeros(batch, nch, n).cuda()
y_i = torch.zeros(batch, nch, n).cuda()
if torch.cuda.is_available():
x1_r, x1_i = cfft.fft(x, x_i)
y1_r, y1_i = cfft.fft(y, y_i)
print(x1_i)
x0, z0 = cfft.ifft(x1_r * y1_r, x1_i * y1_i)
print(z0)
else:
print("Cuda not available, cannot test.")
def step(self, closure=None):
"""
Performs a single optimization step.
Arguments:
closure (callable, optional): A closure that reevaluates the model and returns the loss.
"""
loss = None
if closure is not None:
loss = closure()
for group in self.param_groups:
weight_decay = group['weight_decay']
momentum = group['momentum']
dampening = group['dampening']
nesterov = group['nesterov']
# Update the parameters
idx = 0
for param in group['params']:
if self.sizes[idx] > 2:
#if idx < 1:
if param.grad is None:
continue
tmp = param.grad.view(-1, self.sizes[idx])
# The standard deviation of the injected noise
eps = math.sqrt(2. * self.lr)
tmp1 = tmp.data
tmp2 = eps * torch.randn(tmp.shape).cuda()
re, im = fft.fft(tmp1, self.zero_Ns[idx])
re, im = re * self.coeffs[idx], im * self.coeffs[idx]
tmp1 = fft.ifft(re, im)[0]
re, im = fft.fft(tmp2, self.zero_Ns[idx])
re, im = re * self.coeffs_noise[
idx], im * self.coeffs_noise[idx]
tmp2 = fft.ifft(re, im)[0]
tmp = tmp1 + tmp2
tmp = tmp.view(param.grad.size())
param.grad.data = tmp
#print(p.grad)
idx += 1
d_p = param.grad.data
if weight_decay != 0:
d_p.add_(weight_decay, param.data)
if momentum != 0:
param_state = self.state[param]
if 'momentum_buffer' not in param_state:
buf = param_state[
'momentum_buffer'] = torch.zeros_like(
param.data)
buf.mul_(momentum).add_(d_p)
else:
buf = param_state['momentum_buffer']
buf.mul_(momentum).add_(1 - dampening, d_p)
if nesterov:
d_p = d_p.add(momentum, buf)
else:
d_p = buf
param.data.add_(-group['lr'], d_p)
return loss
def compact_bilinear_pooling_layer(bottom1, bottom2, output_dim, not_variable=True, sum_pool=False,
rand_h_1=None, rand_s_1=None, rand_h_2=None, rand_s_2=None,
seed_h_1=1, seed_s_1=3, seed_h_2=5, seed_s_2=7, sequential=True,
compute_size=128):
"""
Compute compact bilinear pooling over two bottom inputs. Reference:
Yang Gao, et al. "Compact Bilinear Pooling." in Proceedings of IEEE
Conference on Computer Vision and Pattern Recognition (2016).
Akira Fukui, et al. "Multimodal Compact Bilinear Pooling for Visual Question
Answering and Visual Grounding." arXiv preprint arXiv:1606.01847 (2016).
Args:
bottom1: 1st input, 4D Tensor of shape [batch_size, height, width, input_dim1].
bottom2: 2nd input, 4D Tensor of shape [batch_size, height, width, input_dim2].
output_dim: output dimension for compact bilinear pooling.
sum_pool: (Optional) If True, sum the output along height and width
dimensions and return output shape [batch_size, output_dim].
Otherwise return [batch_size, height, width, output_dim].
Default: True.
rand_h_1: (Optional) an 1D numpy array containing indices in interval
`[0, output_dim)`. Automatically generated from `seed_h_1`
if is None.
rand_s_1: (Optional) an 1D numpy array of 1 and -1, having the same shape
as `rand_h_1`. Automatically generated from `seed_s_1` if is
None.
rand_h_2: (Optional) an 1D numpy array containing indices in interval
`[0, output_dim)`. Automatically generated from `seed_h_2`
if is None.
rand_s_2: (Optional) an 1D numpy array of 1 and -1, having the same shape
as `rand_h_2`. Automatically generated from `seed_s_2` if is
None.
sequential: (Optional) if True, use the sequential FFT and IFFT
instead of tf.batch_fft or tf.batch_ifft to avoid
out-of-memory (OOM) error.
Note: sequential FFT and IFFT are only available on GPU
Default: True.
compute_size: (Optional) The maximum size of sub-batch to be forwarded
through FFT or IFFT in one time. Large compute_size may
be faster but can cause OOM and FFT failure. This
parameter is only effective when sequential == True.
Default: 128.
Returns:
Compact bilinear pooled results of shape [batch_size, output_dim] or
[batch_size, height, width, output_dim], depending on `sum_pool`.
"""
# Static shapes are needed to construction count sketch matrix
input_dim1 = bottom1.size()[-1]
input_dim2 = bottom2.size()[-1]
# Step 0: Generate vectors and sketch matrix for tensor count sketch
# This is only done once during graph construction, and fixed during each
# operation
if rand_h_1 is None:
np.random.seed(seed_h_1)
rand_h_1 = np.random.randint(output_dim, size=input_dim1)
if rand_s_1 is None:
np.random.seed(seed_s_1)
rand_s_1 = 2*np.random.randint(2, size=input_dim1) - 1
sparse_sketch_matrix1 = _generate_sketch_matrix_pyt(rand_h_1, rand_s_1, output_dim)
if rand_h_2 is None:
np.random.seed(seed_h_2)
rand_h_2 = np.random.randint(output_dim, size=input_dim2)
if rand_s_2 is None:
np.random.seed(seed_s_2)
rand_s_2 = 2*np.random.randint(2, size=input_dim2) - 1
sparse_sketch_matrix2 = _generate_sketch_matrix_pyt(rand_h_2, rand_s_2, output_dim)
# Step 1: Flatten the input tensors and count sketch
bottom1_flat = bottom1_pyt.view(-1, input_dim1).float()
bottom2_flat = bottom2_pyt.view(-1, input_dim2).float()
# Essentially:_
# sketch1 = bottom1 * sparse_sketch_matrix
# sketch2 = bottom2 * sparse_sketch_matrix
# But tensorflow only supports left multiplying a sparse matrix, so:
# sketch1 = (sparse_sketch_matrix.T * bottom1.T).T
# sketch2 = (sparse_sketch_matrix.T * bottom2.T).T
if not_variable == True:
sketch1 = T.mm(sparse_sketch_matrix1.t().cuda(), bottom1_flat.t()).t()
sketch2 = T.mm(sparse_sketch_matrix2.t().cuda(), bottom2_flat.t()).t()
else:
dense1 = Variable(sparse_sketch_matrix1.to_dense(), requires_grad = True).cuda()
dense2 = Variable(sparse_sketch_matrix2.to_dense(), requires_grad = True).cuda()
sketch1 = T.matmul(bottom1_flat, dense1).cuda()
sketch2 = T.matmul(bottom2_flat, dense2).cuda()
# Step 2: FFT
if not_variable == True:
fft1_real, fft1_img = fft.fft(sketch1, T.zeros(sketch1.size()).cuda())
fft2_real, fft2_img = fft.fft(sketch2, T.zeros(sketch2.size()).cuda())
else:
f = Fft.Fft()
fft1_real, fft1_img = f(sketch1, Variable(T.zeros(sketch1.size())).cuda())
fft2_real, fft2_img = f(sketch2, Variable(T.zeros(sketch2.size())).cuda())
# Step 3: Elementwise product
fft_product_real = fft1_real * fft2_real - fft1_img * fft2_img # The result of only real number part
fft_product_img = fft1_real * fft2_img + fft2_real * fft1_img # The result of only real number part
# Step 4: Inverse FFT and reshape back
# Compute output shape dynamically: [batch_size, height, width, output_dim]
#cbp_flat = tf.real(_ifft(fft_product, sequential, compute_size))
if not_variable == True:
cbp_flat, _ = fft.ifft(fft_product_real, fft_product_img)
else:
fi = Fft.Ifft()
cbp_flat, _ = fi(fft_product_real, fft_product_img)
output_shape = T.Size([bottom1.size()[0], bottom1.size()[1],bottom1.size()[2], output_dim])
cbp = cbp_flat.view(output_shape)
# Step 5: Sum pool over spatial dimensions, if specified
if sum_pool:
cbp = T.sum(T.sum(cbp, dim=1),dim=2)
return cbp