def __init__(self,
params,
lr=required,
sigma=1.,
momentum=0.,
dampening=0.,
weight_decay=0.,
nesterov=False,
vecs=None):
defaults = dict(lr=lr,
sigma=sigma,
momentum=momentum,
dampening=dampening,
weight_decay=weight_decay,
nesterov=nesterov,
vecs=vecs)
if nesterov and (momentum <= 0 or dampening != 0):
raise ValueError(
"Nesterov momentum requires a momentum and zero dampening")
super(LSSGLD2, self).__init__(params, defaults)
coeffs_noise = []
for vec in vecs:
c = torch.Tensor(vec).cuda()
zero_N = torch.zeros(1, vec.shape[1]).cuda()
c_fft, _ = fft.fft(c, zero_N)
coeff = 1. / c_fft
coeffs_noise.append(coeff)
sizes = []
for param in self.param_groups[0]['params']:
sizes.append(torch.numel(param))
coeffs = []
zero_Ns = []
for size in sizes:
if size > 2:
c = np.zeros(shape=(1, size))
c[0, 0] = -2.
c[0, 1] = 1.
c[0, -1] = 1.
c = torch.Tensor(c).cuda()
zero_N = torch.zeros(1, size).cuda()
c_fft, _ = fft.fft(c, zero_N)
coeff = 1. / (1. - sigma * c_fft)
coeffs.append(coeff)
zero_Ns.append(zero_N)
self.lr = lr
self.sigma = sigma
self.sizes = sizes
self.coeffs = coeffs
self.zero_Ns = zero_Ns
self.coeffs_noise = coeffs_noise
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 test_fft_cuFFT():
"""Next part is for comparison between this FFT and cuFFT pytorch-fft https://github.com/locuslab/pytorch_fft"""
import pytorch_fft.fft as pyfft
import time
# A = time domain data B = frequency domain data
A_real, A_imag = torch.randn(4, 2048, 1024).cuda().double(), torch.zeros(
4, 2048, 1024).cuda().double()
start_time = time.clock()
B_real, B_imag = pyfft.fft(A_real, A_imag)
py_fft_time = time.clock() - start_time
x = Variable(A_real) # my FFT takes a Variable directly, not a tensor
start_time = time.clock()
my_B_real, my_B_imag = fft.FFT_torch(x)
my_fft_time = time.clock() - start_time
B_real = Variable(B_real.double().cuda())
B_imag = Variable(B_imag.double().cuda())
print("FFT_torch is equal to pytorch-fft? : ",
np.allclose(my_B_real.cpu().data.numpy(),
B_real.cpu().data.numpy()))
print("\nTime pytorch-fft: {:.3f}ms ".format((py_fft_time) * 1000))
print("\nTime my FFT_torch: {:.3f}ms".format((my_fft_time) * 1000))
def __init__(self, params, lr=required, sigma=0.05, momentum=0, dampening=0, weight_decay=0, nesterov=False):
defaults = dict(lr=lr, sigma=sigma, momentum=momentum, dampening=dampening, weight_decay=weight_decay, nesterov=nesterov)
if nesterov and (momentum <= 0 or dampening != 0):
raise ValueError("Nesterov momentum requires a momentum and zero dampening")
super(Grad_SJO_SGD, self).__init__(params, defaults)
sizes = []
for param in self.param_groups[0]['params']:
sizes.append(torch.numel(param))
coeffs = []
zero_Ns = []
for size in sizes:
c = np.zeros(shape=(1, size))
c[0, 0] = -2.
c[0, 1] = 1.
c[0, -1] = 1.
c = torch.Tensor(c).cuda()
zero_N = torch.zeros(1, size).cuda()
c_fft, _ = fft.fft(c, zero_N)
coeff = 1. / (1.-sigma*c_fft)
coeffs.append(coeff)
zero_Ns.append(zero_N)
self.sigma = sigma
self.sizes = sizes
self.coeffs = coeffs
self.zero_Ns = zero_Ns
def forward(self, X_re, X_im):
X_re, X_im = make_contiguous(X_re, X_im)
k_re, k_im = fft(X_re, X_im)
if self.norm == 'ortho':
N = np.sqrt(k_re.size(-1))
k_re /= N
k_im /= N
return k_re, k_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 = fft(grad_output_im, grad_output_re)
if self.norm == 'ortho':
N = np.sqrt(gi.size(-1))
gi /= N
gr /= N
return gr, gi
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 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 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 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