Hi,
I am trying to fine-tune the last two convolutional blocks of the VGG 16 model. I have added a conv2d layer with 512 filters on top of VGG 16, for visualization of the output. The data for fine-tuning has both RGB images and binary masks. The code for training is like this:
device = t.device(‘cuda:0’ if t.cuda.is_available() else ‘cpu’)
Here is the loss and optimizer definition
criterion = nn.NLLLoss()
optimizer = t.optim.Adam(model.parameters(), 5e-4, (0.9, 0.999), eps=1e-08, weight_decay=1e-4)
start_epoch = 1
steps_loss = 50
my_start_time = time.time()
The training loop
total_steps = len(train_loader)
epochs = 30
print(f"{epochs} epochs, {total_steps} total_steps per epoch")
for epoch in range(epochs):
print(" — — — TRAINING — EPOCH", epoch, " — — -")
epoch_loss = []
time_train = []
model.train()
for i, (images, masks) in enumerate(train_loader, 1):
start_time = time.time()
images = images.to(device)
masks = masks.type(t.LongTensor)
masks = masks.reshape(masks.shape[0], masks.shape[2], masks.shape[3])
masks = masks.to(device)
    # Forward pass
    output1 = model(images)
    softmax = F.log_softmax(outputs1, dim=1)
    loss = criterion(softmax, masks)
    
    # Backward and optimize
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    epoch_loss.append(loss.item())
    time_train.append(time.time()-start_time)
    
    if steps_loss > 0 and i%steps_loss == 0:
        average = sum(epoch_loss)/len(epoch_loss)
        print('loss: {average:0.4} (epoch: {epoch}, step: {i})', "//Avg time/img: %.4f s" % (sum(time_train)/len(time_train)/batch_size))

        average_epoch_loss_train = sum(epoch_loss)/len(epoch_loss)

It is throwing a runtime error like
RuntimeError                              Traceback (most recent call last)
 in 
29         output1 = model(images)
30         softmax = F.log_softmax(outputs1, dim=1)
—> 31         loss = criterion(softmax, masks)
32
33         # Backward and optimize
~/yes/lib/python3.7/site-packages/torch/nn/modules/module.py in call(self, *input, **kwargs)
530             result = self._slow_forward(*input, **kwargs)
531         else:
–> 532             result = self.forward(*input, **kwargs)
533         for hook in self._forward_hooks.values():
534             hook_result = hook(self, input, result)
~/yes/lib/python3.7/site-packages/torch/nn/modules/loss.py in forward(self, input, target)
202
203     def forward(self, input, target):
–> 204         return F.nll_loss(input, target, weight=self.weight, ignore_index=self.ignore_index, reduction=self.reduction)
205
206
~/yes/lib/python3.7/site-packages/torch/nn/functional.py in nll_loss(input, target, weight, size_average, ignore_index, reduce, reduction)
1838         ret = torch._C._nn.nll_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
1839     elif dim == 4:
-> 1840         ret = torch._C._nn.nll_loss2d(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
1841     else:
1842         # dim == 3 or dim > 4
RuntimeError: input and target batch or spatial sizes don’t match: target [8 x 360 x 640], input [8 x 512 x 22 x 40] at /opt/conda/conda-bld/pytorch_1579022060824/work/aten/src/THCUNN/generic/SpatialClassNLLCriterion.cu:23
Please help me. I am very new to Pytorch. Sorry if my question sounds dumb. Thank you for the help in advance.