Hi everyone,
I’m trying to use nn.DataParallel to have a multi-gpu training, but I encountered the following error. I had a look at the various threads, but I wasn’t able to fix the issue:
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)
I’ve used the following lines before importing torch to limit the visible GPUs (as I do for the training on a single gpu):
os.environ["CUDA_VISIBLE_DEVICES"] = "0, 2, 3"
os.getcwd()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
The model it’s used in this way (with weights from another training):
model = ResNet()
model.load_state_dict(torch.load('rot_weights.pt'))
model = nn.DataParallel(model, device_ids=[0, 2, 3]).to(device)
optimizer = optim.Adam(model.parameters(),
lr=learning_rate,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=0,
amsgrad=False)
The model is a ResNet-18 from which I want just the feature extraction part, modified for my regression problem:
class ResNet(nn.Module):
def __init__(self):
super(ResNet, self).__init__()
self.model = pretrainedmodels.__dict__['resnet18'](pretrained='imagenet')
self.regression_layer = nn.Sequential(nn.Linear(512, 6))
def forward(self, x):
batch_size ,_,_,_ = x.shape #taking out batch_size from input image
x = self.model.features(x)
x = torch.nn.functional.adaptive_avg_pool2d(x,1).reshape(batch_size,-1) # then reshaping the batch_size
x = self.regression_layer(x)
x = compute_rotation_matrix_from_ortho6d(x.view(batch_size, -1))
return x
def compute_rotation_matrix_l2_loss(self, gt_rotation_matrix, predict_rotation_matrix):
loss_function = nn.MSELoss()
loss = loss_function(predict_rotation_matrix, gt_rotation_matrix)
return loss
def compute_rotation_matrix_geodesic_loss(self, gt_rotation_matrix, predict_rotation_matrix):
theta = compute_geodesic_distance_from_two_matrices(gt_rotation_matrix, predict_rotation_matrix)
error = theta.mean()
return error
Any suggestion would be really appreciated!