My problem is that after transposing tensor two times its gradient disappears.
Since the model’s weight matrix is large, I performed matrix multiplication as output = weight.mm(input.t()).t() instead of output=input.mm(weight.t())
However, it makes weight's gradient to disappear.
Using output=input.mm(weight.t()) makes the model works fine.
Could you find out what is wrong?
Below is my code
import numpy as np
import torch
import torch.nn as nn
import torch.nn.init as init
import torch.optim as optim
device_cpu=torch.device("cpu")
device_cuda_list=[torch.device("cuda:{}".format(i)) for i in range(6)]
import math
import torch
import torch.nn as nn
class CustomizedLinearFunction(torch.autograd.Function):
"""
autograd function which masks it's weights by 'mask'.
"""
# Note that both forward and backward are @staticmethods
@staticmethod
# bias, mask is an optional argument
def forward(ctx, input, weight, bias=None, mask=None):
if mask is not None:
# change weight to 0 where mask == 0
weight.mul_(mask)
output = weight.mm(input.t()).t()
#output=input.mm(weight.t())
if bias is not None:
output += bias.unsqueeze(0).expand_as(output)
ctx.save_for_backward(input, weight, bias, mask)
return output
# This function has only a single output, so it gets only one gradient
@staticmethod
def backward(ctx, grad_output):
input, weight, bias, mask = ctx.saved_tensors
grad_input = grad_weight = grad_bias = grad_mask = None
if ctx.needs_input_grad[0]:
grad_input = grad_output.mm(weight)
if ctx.needs_input_grad[1]:
grad_weight = grad_output.t().mm(input)
if mask is not None:
# change grad_weight to 0 where mask == 0
#grad_weight=grad_weight*mask
grad_weight.mul_(mask) #*
#if bias is not None and ctx.needs_input_grad[2]:
if ctx.needs_input_grad[2]:
grad_bias = grad_output.sum(0).squeeze(0)
return grad_input, grad_weight, grad_bias, grad_mask
class CustomizedLinear(nn.Module):
def __init__(self, mask, bias=True):
super(CustomizedLinear, self).__init__()
self.input_features = mask.shape[1]
self.output_features = mask.shape[0]
if isinstance(mask, torch.Tensor):
self.mask = mask.type(torch.float)
else:
self.mask = torch.tensor(mask, dtype=torch.float)
self.mask = nn.Parameter(self.mask, requires_grad=False)
self.weight = nn.Parameter(torch.Tensor(self.output_features, self.input_features))
if bias:
self.bias = nn.Parameter(torch.Tensor(self.output_features))
else:
# You should always register all possible parameters, but the
# optional ones can be None if you want.
self.register_parameter('bias', None)
self.reset_parameters()
# mask weight
self.weight.dat = self.weight.data.mul_(self.mask)
#self.weight.mul_(self.mask)
def reset_parameters(self):
init.kaiming_uniform_(self.weight, a=math.sqrt(5))
if self.bias is not None:
fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
bound = 1 / math.sqrt(fan_in)
init.uniform_(self.bias, -bound, bound)
def forward(self, input):
# See the autograd section for explanation of what happens here.
return CustomizedLinearFunction.apply(input, self.weight, self.bias, self.mask)
def extra_repr(self):
# (Optional)Set the extra information about this module. You can test
# it by printing an object of this class.
return 'input_features={}, output_features={}, bias={}'.format(
self.input_features, self.output_features, self.bias is not None
)
adjacency=(np.random.rand(100,50)>0.5).astype(int)
adjacency
model=nn.Sequential(CustomizedLinear(mask=adjacency,bias=True)
,nn.ELU()
).to(device_cuda_list[-2])
#optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9,weight_decay=5e-07)
optimizer = optim.Adagrad(model.parameters(),weight_decay=5e-07)
#scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=5,verbose=1)
loss_function = nn.BCEWithLogitsLoss()
target_model=model[0]
print(target_model.weight,target_model.weight.shape)
print(target_model.weight.grad)
print(target_model.bias,target_model.bias.shape)
print(target_model.bias.grad)
#sample=torch.rand(19000).to(device_cuda_list[0])
input_sample=torch.rand([2,50]).to(device_cuda_list[-2])
out_sample=model(input_sample)
#out_sample
loss=loss_function(out_sample,torch.zeros([2,100]).to(device_cuda_list[-2]))
loss.backward()
optimizer.step()
print(target_model.weight,target_model.weight.shape)
print(target_model.weight.grad,target_model.weight.grad.shape)
print(target_model.bias,target_model.bias.shape)
print(target_model.bias.grad,target_model.bias.grad.shape)
---------------------------------------------------------------------------
AttributeError Traceback (most recent call last)
<ipython-input-20-24264289d036> in <module>
----> 1 target_model.weight.grad,target_model.weight.grad.shape
AttributeError: 'NoneType' object has no attribute 'shape'