Error with DataParallel

username: dmandair
cooked: <p>Hi everyone - really would appreciate your help on this. I’ve been trying to speed up my resnet (basically resnet34, I customized it to expand the features by adding an additional block at the end). To this end, I added nn.DataParallel as a wrapper for my model. The code below worked entirely fine prior to the additions of nn.DataParallel (and using model.cuda() and inputs.cuda() vs to(device) which I used prior). The first run of this model, which is done with only the final output layer and batchnorm layers unfrozen, works fine. I then take this model and unfreeze all weights in all layers of the network, and run this model through the training loop once again. This produces the following error. I’m really not sure how to fix this, although I’m guessing it likely has to do with using the previously trained model and putting that through another training loop (inside of which I again call model.cuda()). I’ve pasted as much of my code as I thought would be relevant, including the training loop and calls to partial unfreezing. Would really appreciate your help and thanks again!</p>
<p>Traceback (most recent call last):<br>
File “/wynton/protected/home/ichs/dmandair/BRCA/train_val.py”, line 395, in <br>
trained_model = train_model(trained_model, criterion, optimizer_ft, exp_lr_scheduler, dataloader_dict, device, NUM_EPOCHS)<br>
File “/wynton/protected/home/ichs/dmandair/BRCA/train_val.py”, line 155, in train_model<br>
outputs = model(inputs)<br>
File “/wynton/protected/home/ichs/dmandair/anaconda3/envs/ONC_EXP/lib/python3.9/site-packages/torch/nn/modules/module.py”, line 1102, in _call_impl<br>
return forward_call(*input, **kwargs)<br>
File “/wynton/protected/home/ichs/dmandair/anaconda3/envs/ONC_EXP/lib/python3.<br>
9/site-packages/torch/nn/parallel/data_parallel.py”, line 168, in forward<br>
outputs = self.parallel_apply(replicas, inputs, kwargs)<br>
File “/wynton/protected/home/ichs/dmandair/anaconda3/envs/ONC_EXP/lib/python3.9/site-packages/torch/nn/parallel/data_parallel.py”, line 178, in parallel_apply<br>
return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])<br>
File “/wynton/protected/home/ichs/dmandair/anaconda3/envs/ONC_EXP/lib/python3.9/site-packages/torch/nn/parallel/parallel_apply.py”, line 86, in parallel_apply<br>
output.reraise()<br>
File “/wynton/protected/home/ichs/dmandair/anaconda3/envs/ONC_EXP/lib/python3.9/site-packages/torch/_utils.py”, line 434, in reraise<br>
raise exception<br>
RuntimeError: Caught RuntimeError in replica 0 on device 0.<br>
Original Traceback (most recent call last):<br>
File “/wynton/protected/home/ichs/dmandair/anaconda3/envs/ONC_EXP/lib/python3.9/site-packages/torch/nn/parallel/parallel_apply.py”, line 61, in _worker<br>
output = module(*input, **kwargs)<br>
File “/wynton/protected/home/ichs/dmandair/anaconda3/envs/ONC_EXP/lib/python3.9/site-packages/torch/nn/modules/module.py”, line 1102, in _call_impl<br>
return forward_call(*input, **kwargs)<br>
File “/wynton/protected/home/ichs/dmandair/anaconda3/envs/ONC_EXP/lib/python3.9/site-packages/torch/nn/parallel/data_parallel.py”, line 168, in forward<br>
outputs = self.parallel_apply(replicas, inputs, kwargs)<br>
File “/wynton/protected/home/ichs/dmandair/anaconda3/envs/ONC_EXP/lib/python3.<br>
9/site-packages/torch/nn/parallel/data_parallel.py”, line 178, in parallel_apply<br>
return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])<br>
File “/wynton/protected/home/ichs/dmandair/anaconda3/envs/ONC_EXP/lib/python3.9/site-packages/torch/nn/parallel/parallel_apply.py”, line 86, in parallel_apply<br>
output.reraise()<br>
File “/wynton/protected/home/ichs/dmandair/anaconda3/envs/ONC_EXP/lib/python3.9/site-packages/torch/_utils.py”, line 434, in reraise<br>
raise exception<br>
RuntimeError: Caught RuntimeError in replica 1 on device 1.<br>
Original Traceback (most recent call last):<br>
File “/wynton/protected/home/ichs/dmandair/anaconda3/envs/ONC_EXP/lib/python3.9/site-packages/torch/nn/parallel/parallel_apply.py”, line 61, in _worker<br>
output = module(*input, **kwargs)<br>
File “/wynton/protected/home/ichs/dmandair/anaconda3/envs/ONC_EXP/lib/python3.9/site-packages/torch/nn/modules/module.py”, line 1102, in _call_impl<br>
return forward_call(*input, **kwargs)<br>
File “/wynton/protected/home/ichs/dmandair/BRCA/train_val.py”, line 222, in forward<br>
x = self.base(x)<br>
File “/wynton/protected/home/ichs/dmandair/anaconda3/envs/ONC_EXP/lib/python3.9/site-packages/torch/nn/modules/module.py”, line 1102, in _call_impl<br>
return forward_call(*input, **kwargs)<br>
File “/wynton/protected/home/ichs/dmandair/anaconda3/envs/ONC_EXP/lib/python3.<br>
9/site-packages/torch/nn/modules/container.py”, line 141, in forward<br>
input = module(input)<br>
File “/wynton/protected/home/ichs/dmandair/anaconda3/envs/ONC_EXP/lib/python3.9/site-packages/torch/nn/modules/module.py”, line 1102, in _call_impl<br>
return forward_call(*input, **kwargs)<br>
File “/wynton/protected/home/ichs/dmandair/anaconda3/envs/ONC_EXP/lib/python3.9/site-packages/torch/nn/modules/conv.py”, line 446, in forward<br>
return self._conv_forward(input, self.weight, self.bias)<br>
File “/wynton/protected/home/ichs/dmandair/anaconda3/envs/ONC_EXP/lib/python3.9/site-packages/torch/nn/modules/conv.py”, line 442, in _conv_forward<br>
return F.conv2d(input, weight, bias, self.stride,<br>
<strong>RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cuda:0! (when checking argument for argument weight in method wrapper__cudnn_convolution)</strong></p>
<pre><code>def train_model(model, criterion, optimizer, scheduler, dataloaders, device, num_epochs=4):
     since = time.time()
     model = nn.DataParallel(model, device_ids=[0, 1, 2, 3])
     model = model.cuda()
     best_model_wts = copy.deepcopy(model.state_dict())
     best_acc = 0.0
     for epoch in range(num_epochs):
          print('Epoch {}/{}'.format(epoch, num_epochs - 1))
          print('-' * 10)
    # Each epoch has a training and validation phase
        for phase in ['train', 'val']:
             if phase == 'train':
                 model.train()  # Set model to training mode
             else:
                 model.eval()  # Set model to evaluate mode

            running_loss = 0.0
            running_corrects = 0

            # Iterate over data.
            for inputs, labels, _, _ in dataloaders[phase]:
                inputs = inputs.cuda()
                labels = labels.cuda()
                # zero the parameter gradients
                optimizer.zero_grad()
            # forward
            # track history if only in train
                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(inputs)
                    _, preds = torch.max(outputs, 1)
                    loss = criterion(outputs, labels)

                    # backward + optimize only if in training phase
                    if phase == 'train':
                        loss.backward()
                        torch.nn.utils.clip_grad_norm_(model.parameters(), MAX_NORM)
                        optimizer.step()

                # statistics
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)
            if phase == 'train':
                scheduler.step()

                epoch_loss = running_loss / len(train_tile_ds)
                epoch_acc = running_corrects.double() / len(train_tile_ds)
            else:
                epoch_loss = running_loss / len(val_tile_ds)
                epoch_acc = running_corrects.double() / len(val_tile_ds)

            print('{} Loss: {:.4f} Acc: {:.4f}'.format(
                phase, epoch_loss, epoch_acc))

            # deep copy the model
            if phase == 'val' and epoch_acc &gt; best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())

        print()

    time_elapsed = time.time() - since
    print('Training complete in {:.0f}m {:.0f}s'.format(
        time_elapsed // 60, time_elapsed % 60))
    print('Best val Acc: {:4f}'.format(best_acc))

    # load best model weights
    model.load_state_dict(best_model_wts)
    return model

def get_new_base(orig_model):
    return nn.Sequential(*list(orig_model.children())[:-2])


def create_new_block(in_channel, out_channel):
    conv_layer = nn.Conv2d(in_channel, out_channel, kernel_size=(1, 1), stride=(1, 1), bias=False)
    batch_layer = nn.BatchNorm2d(out_channel, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    relu = nn.ReLU(inplace=True)
    average_pool = nn.AdaptiveAvgPool2d(output_size=(1, 1))
    return [conv_layer, batch_layer, relu, average_pool]

class CustomResNet(nn.Module):
    def __init__(self, out_c, out_feat):
        super().__init__()
        base = models.resnet34(pretrained=True)

        base_channels_op = base.fc.in_features
        self.base = get_new_base(base)
        self.new_layers = nn.Sequential(*create_new_block(base_channels_op, out_c))
        self.final = nn.Linear(out_c, out_feat)

    def forward(self, x):
        x = self.base(x)
        x = self.new_layers(x)
        x = torch.flatten(x,1)
        return self.final(x)

def set_grad(m, b):
    if isinstance(m, (nn.Linear,nn.BatchNorm2d)): return
    if hasattr(m, 'weight'):
        for p in m.parameters(): p.requires_grad_(b)


train_loader = DataLoader(train_tile_ds, batch_size=BS,
                                       shuffle = True, num_workers=4)
validation_loader = DataLoader(val_tile_ds, batch_size=BS, shuffle = False, num_workers=4)

dataloader_dict = {'train': train_loader, 'val': validation_loader}

mod_resnet = CustomResNet(768, 2)

mod_resnet.apply(partial(set_grad, b=False))

criterion = nn.CrossEntropyLoss()

optimizer_ft = optim.SGD([{'params': mod_resnet.base.parameters(), 'lr': 1e-2},
            {'params': mod_resnet.new_layers.parameters(), 'lr': 1e-3},
            {'params': mod_resnet.final.parameters(), 'lr': 1e-3}
        ], momentum=0.9)

exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)

trained_model = train_model(mod_resnet, criterion, optimizer_ft, exp_lr_scheduler, dataloader_dict, device, NUM_EPOCHS)

optimizer_ft = optim.SGD([{'params': trained_model.module.base.parameters(), 'lr': 1e-3},
            {'params': trained_model.module.new_layers.parameters(), 'lr': 1e-3},
            {'params': trained_model.module.final.parameters(), 'lr': 1e-3}
        ], momentum=0.9)

trained_model = train_model(trained_model, criterion, optimizer_ft, exp_lr_scheduler, dataloader_dict, device, NUM_EPOCHS)</code></pre>

username: ptrblck
cooked: <p>Check the <code>.device</code> attribute of the model parameters as well as the input data in the <code>forward</code> method of your model and try to isolate which tensor creates the device mismatch.<br>
Once this is done, check how this tensor/parameter is created and if you are manually pushing this tensor to a specific device instead of letting <code>nn.DataParallel</code> handle it.</p>