Hi,
I’m trying to to train the fcn_resnet101 model with my own data to do image semantic segmentation. I’m trying do this implement this by trying to use the fine tuning tutorial, structure. Right now I am able to train, but I’m pretty sure I’m not using the right loss functions, and optimizers and probably a lot more missing functions. I am just beginning with Pytorch fyi.
My images come in batch shape: (N, 3, 256, 256)
My mask come in batch shape: (N, 256, 256) # each pixel is a class 0,1 or 2
I changed the last classifier layer of the to my 3 amount of classes (bg, window frame, glass), like this ( I hope this is the correct way):
num_classes = 3  # (background, kozine, glass)
model.classifier[4] = nn.Conv2d(512, num_classes, kernel_size=(1,1), stride=(1,1))

These are following model params I have put in the model:
# model params
num_epochs = 3
criterion = nn.modules.loss.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.005, momentum=0.9, weight_decay=0.0005)
dataloader = DataLoader(trainset, batch_size=5)

The model training code itself.
start_time = time.time()
    
best_model = copy.deepcopy(model.state_dict())
best_acc = 0.0
val_acc_hist = []

for epoch in range(num_epochs):
    print('Epoch: {}/ {}'.format(epoch + 1, num_epochs))

    epoch_loss = 0    
    running_loss = 0.0
    running_corrects = 0

    # Iterate batches over dataloader
    for k, (inputs, labels) in enumerate(dataloader):
        print(f'processing batch {k}...')
        
        # move to gpu 
        _inputs = inputs.cuda()
        _labels = labels.cuda()

        # zero the gradients again
        optimizer.zero_grad()

        # forward
        outputs = model(_inputs) # one batch forward 
        loss1 = criterion(outputs['out'], _labels.long()) # get score of output
        loss2 = criterion(outputs['aux'], _labels.long())
        loss = loss1 + 0.4*loss2
        loss.backward()
        optimizer.step() # move weights

        # save preds and take highest prediction
        _, preds = torch.max(outputs['out'], 1)

        running_loss += loss.item() * inputs.size(0) # length of the batch
        running_corrects += torch.sum(preds == _labels.data) # amount of correct pixels in batch

    # epoch statistics
    epoch_loss = running_loss / len(dataloader.dataset) # avg loss of whole dataset
    epoch_acc = (running_corrects.double() / len(dataloader.dataset)) / 262144 # avg correct pixels

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

    # evaluation phase
    if epoch_acc > best_acc:
        best_acc = epoch_acc
        best_model_wts = copy.deepcopy(model.state_dict())

    val_acc_hist.append(epoch_acc)


# after finishing training
time_elapsed = time.time() - start_time
print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
print('Best val acc: {:4f}'.format(best_acc))

test_model = model.load_state_dict(best_model_wts)

The loss, and accuracy is just stays the same for most part, I definitely am doing it wrong.
epoch 1 - Loss: 0.0056 Acc: 0.0823
epoch 2 - Loss: 0.0048 Acc: 0.0823
epoch 3 - Loss: 0.0042 Acc: 0.0823

Can anybody point where to go from here to fix it? Or is there an official implementation where I can look it like the mask rcnn tutorial?
Any help is appreciated!