Hello! I would like to compute the gradients of a tensor matrix with respect to another set of gradients. Below code works fine on a single GPU. However, when multiple GPUs are used, the error occurs. Can anyone share with me how to handle this? Thanks!
May I ask if you could have a look? Thank you very much! @ptrblck @albanD
My code:
#model
model = models.resnet50(pretrained=True)
model.fc = nn.Linear(2048, num_classes)    

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if torch.cuda.device_count() > 1:
    model = nn.DataParallel(model)
model = model.to(device) 

#T is a tensor, which I would like to compute its gradients based on another set of gradients
T = torch.from_numpy(numpy.random.rand(num_classes, num_classes)).float()
T = T.requires_grad_(True).to(device)

#do not want to update the weights of the original model, so get another copy of the model
net = copy.deepcopy(model) 

#copy net's weights for later gradients computation manually
original_weights = OrderedDict()  
for param, weights in net.named_parameters():
    if not weights.requires_grad:
        print(param)   
    else:
        original_weights[param] = copy.deepcopy(weights)        
original_weights_keys = tuple(original_weights.keys())

#data_s is images, target_s is labels
with torch.enable_grad():
    logits = net(data_s)     
    #pre1 is the (target_s)-th row of the tensor matrix T, basically, loss_ is crossentropy loss
    pre1 = T[torch.cuda.LongTensor(target_s.data)]
    pre2 = torch.mul(F.softmax(logits, dim=1), pre1)  
    loss_ = -(torch.log(pre2.sum(1))).sum(0)/float(len(target_s))
    
    #manually compute gradients
    net_grads = torch.autograd.grad(loss_, net.parameters(), retain_graph=True, create_graph=True, allow_unused=True) 

    #see _del_nested_attr() and _set_nested_attr() in https://discuss.pytorch.org/t/how-to-calculate-2nd-derivative-of-a-likelihood-function/15085/30 
    for param, grad in zip(original_weights_keys, net_grads):                
        if grad is None: 
            continue
        else:
            new_weight = original_weights[param] - learning_rate * grad
            # remove old Parameter
            _del_nested_attr(net, param.split("."))
            # set the Tensor with history
            _set_nested_attr(net, param.split("."), new_weight)
    
    #error occurs from below line, but it is fine when a single GPU is used.
    pre = net(data_g)
    loss = criteria(pre, target_g)
        
    T_grads = torch.autograd.grad(loss, T, allow_unused=True)

with torch.set_grad_enabled(False):    
    T -= learning_rate * T_grads[0]

Error:
File “train_ours_meta_clth1m.py”, line 207, in update_T
pre = net(data_g)
File “/root/miniconda3/envs/pytorch/lib/python3.7/site-packages/torch/nn/modules/module.py”, line 493, in call
result = self.forward(*input, **kwargs)
File “/root/miniconda3/envs/pytorch/lib/python3.7/site-packages/torch/nn/parallel/data_parallel.py”, line 152, in forward
outputs = self.parallel_apply(replicas, inputs, kwargs)
File “/root/miniconda3/envs/pytorch/lib/python3.7/site-packages/torch/nn/parallel/data_parallel.py”, line 162, in parallel_apply
return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
File “/root/miniconda3/envs/pytorch/lib/python3.7/site-packages/torch/nn/parallel/parallel_apply.py”, line 83, in parallel_apply
raise output
File “/root/miniconda3/envs/pytorch/lib/python3.7/site-packages/torch/nn/parallel/parallel_apply.py”, line 59, in _worker
output = module(*input, **kwargs)
File “/root/miniconda3/envs/pytorch/lib/python3.7/site-packages/torch/nn/modules/module.py”, line 493, in call
result = self.forward(*input, **kwargs)
File “/root/miniconda3/envs/pytorch/lib/python3.7/site-packages/torchvision/models/resnet.py”, line 192, in forward
x = self.conv1(x)
File “/root/miniconda3/envs/pytorch/lib/python3.7/site-packages/torch/nn/modules/module.py”, line 493, in call
result = self.forward(*input, **kwargs)
File “/root/miniconda3/envs/pytorch/lib/python3.7/site-packages/torch/nn/modules/conv.py”, line 338, in forward
self.padding, self.dilation, self.groups)
RuntimeError: Expected tensor for argument #1 ‘input’ to have the same device as tensor for argument #2 ‘weight’; but device 1 does not equal 0 (while checking arguments for cudnn_convolution)