I am very new to Pytorch and Deep Learning in general. I have a set of grayscale images which I convert to a 3 channel images by repeating the first channel two more times. I am using the fcn_resnet101 from the pytorch models. The model only predicts one class for all images. Am I missing something from the methodology below.
from loader import DataLoaderSegmentation
import torch.nn as nn
import torch.optim as optim
import numpy as np
from torch.utils.data.sampler import SubsetRandomSampler
batch_size = 1
validation_split = .2
shuffle_dataset = True
n_class = 2
num_epochs = 1
#lr = 1e-4
#momentum = 0
w_decay = 1e-5
step_size = 50
gamma = 0.5
traindata = DataLoaderSegmentation('/home/ubuntu/Downloads/Brain/test0716')
dataset_size = len(traindata)
indices = list(range(dataset_size))
split = int(np.floor(validation_split * dataset_size))
train_indices = indices[split:]
val_indices = indices[:split]
#train_sampler = SubsetRandomSampler(train_indices)
#valid_sampler = SubsetRandomSampler(val_indices)
trainloader = torch.utils.data.DataLoader(traindata, batch_size=batch_size)
validation_loader = torch.utils.data.DataLoader(traindata, batch_size=batch_size)
#trainloader = torch.utils.data.DataLoader(traindata, batch_size=1, shuffle=False, num_workers=2)
model = torchvision.models.segmentation.fcn_resnet101(pretrained=False, progress=True,aux_loss=None, num_classes=2).cuda()
criterion = nn.CrossEntropyLoss().cuda()
#optimizer = optim.RMSprop(model.parameters(), lr=lr, momentum=momentum, weight_decay=w_decay)
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
for iter in range(num_epochs):
for (i,l) in trainloader:
l = l.to(device=device, dtype=torch.int64)
outt = model(i)
loss = criterion(outt['out'], l.squeeze(1))
Usually you would create a train function to train the model for a complete epoch using the training dataset. Afterwards you could run the evaluation on the validation dataset.
Have a look at the ImageNet exmple to see how train and validate are defined.
The same code uses this code to print the progress.
One signal of overfitting is when you observe a decreasing training loss and a static or increasing validation loss.
Further training at this point usually won’t give you any benefits and you would stop at this point.
You would counter it by e.g.
using (or collecting) more data
remove “capacity” of the model, e.g. by reducing the number of parameters
Thank you for your response. I noticed that my Dataset is heavily imbalanced, there are way more pixels representing the backgrounds than the organ (label I am interested in). I have 18000 images of which only 600 have the organ class I am looking for. Is it possible that this is the reason why only one of the classes is learned?
Hi @ptrblck, I tried working with a more even dataset with balanced classes. I also tried playing around with the weight_decay parameter in the optimizer. However, I am still running into the issue where my model is only outputting 1 class. I finally decided to just increase the epochs and overfit the model with training data and perform prediction on a training image. It turns out I am still only getting 1 class when I am supposed to get 2 classes. I tried scouring the internet for a similar issue, but couldn’t find one. Any and all help is very much appreciated.
I have provided a detailed explanation in the following post