Source: Canva

Pytorch Image Segmentation Tutorial For Beginners — II

Making masks for Brain Tumor MRI Images in Pytorch

In the previous tutorial, we prepared data for training. We downloaded the dataset, loaded the images, split the data, defined model structure, downloaded weights, defined training parameters.

In the second part of the tutorial, we train the model and evaluate the results of the model. I

Prepare Training

We first define the device and switch between CPU and GPU when needed.

device = torch.device(‘cuda’) if torch.cuda.is_available() else torch.device(‘cpu’)

We load the model to the device.

model.to(device)

We initialize two lists to store loss values.

loss_history = []
loss_history_val = []

We define best_loss_val, and make it store an infinite value. Thus, best_loss_value acts as an unbounded upper value for comparison. This is useful for finding the lowest value for the loss.

best_loss_val = float(‘inf’)

Train

We start training.

print(“Start train…”)

We make a for loop for every epoch.

We set the train mode.

for epoch in range(50):
    #Train mode
    model.train()
    loss_running = []
    for _, (x,y) in enumerate(train_dataloader):
        x, y = x.float().to(device), y.float().to(device)
        pred = model(x)
        loss = criterion(pred, y)
        loss_running.append(loss.item())
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    loss_history.append(np.mean(loss_running))

We set the evaluation mode.

# Evaluate mode
model.eval()
with torch.no_grad():
loss_val_running = []
for _, (x_val, y_val) in enumerate(val_dataloader):
x_val, y_val = x_val.to(device), y_val.to(device)

We do predictions, calculate loss, append the calculated loss to the list of train and validation losses. (loss_running, loss_val_running)

pred_val = model.forward(x_val) #pred_val = model(x_val)
loss_val= criterion(pred_val, y_val)
loss_val_running.append(loss_val.item())
curr_loss_val = np.mean(loss_val_running)
loss_history_val.append(curr_loss_val)

We save the best weights of the model for later use.

# Save the best weights
if curr_loss_val < best_loss_val:
best_loss_val = curr_loss_val
torch.save(model.state_dict(), ‘/content/drive/My          Drive/Brain-Tumor-Segmentation-Project/best_model.pth’)

We change the learning rate by a learning rate scheduler we defined before.

# Change the learning rate
scheduler.step()

We print the results.

# Print the results
print(“epoch”, epoch, “train loss”, loss_history[-1], “val loss”, loss_history_val[-1])

Here is the block of code.

We load the best model’s weights and evaluate the model.

checkpoint = torch.load('/content/drive/My Drive/Brain-Tumor-Segmentation-Project/best_model.pth', map_location=torch.device('cpu'))
model.load_state_dict(checkpoint)

Let’s now see what we did.

def plot_mask(mask_3d_array, axx): # takes 64*64*64 array
mask_cpu = mask_3d_array.cpu().detach().numpy()
reshaped_mask_cpu = np.reshape(mask_cpu,(512, 512))
print(np.max(reshaped_mask_cpu), np.min(reshaped_mask_cpu))
reshaped_mask_cpu_bin = np.round(reshaped_mask_cpu)
axx.imshow(reshaped_mask_cpu_bin)

Plot some predicted and ground truth masks

dataloader = test_dataloader
ncol = 4
rand_ndx = random.sample(range(0, len(dataloader)), ncol)
fig, ax = plt.subplots(nrows=2,  ncols=ncol, figsize=(20, 10))
i = 0
for n, (x, y) in enumerate(dataloader):
x, y = x.to(device), y.to(device)
if n in rand_ndx:
pred = model.forward(x)
plot_mask(pred[0,0,:,:,:], ax[0][i])
plot_mask(y[0,0,:,:,:], ax[1][i])
i+=1

Predicted(up) and ground truth(down) masks

Congratulations. Here is the link to my project if you want to see more details.

Any suggestions to seymatas@gmail.com will be very appreciated!