I have trained a model fish.pth with 2 channels (fish and background). I want to add a layer so that I train a new model spine.pth which segments 3 classes fish, spine and background. Yes, I can do it in one operation, but I want to learn the method. I made a class which initializes a model with 2 out_channels such that I can load the weights from my pretrained model. Do I have to make a new model or can I just alter model_spine?
Code snippet with what I have tried below. When I get to the training loop outputs = model_spine(inputs) gives a outputs.shape=[20,2,256,256] where I expected outputs.shape=[20,3,256,256]
# 3) --------- Model for spine segmentation ---------------
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
class SpineSegmentationModel2(nn.Module):
def __init__(self):
super(SpineSegmentationModel2, self).__init__()
self.unet = UNet(
spatial_dims=2,
in_channels=1, # 1-channel (grayscale input)
out_channels=2, # 2 classes (background, fish)
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
)
# Load pretrained weights
self.unet.load_state_dict(torch.load(pretrain_path), strict=False)
# Adjust to 3 output layers
self.unet.out_conv = nn.Conv2d(256, 3, kernel_size=(1, 1), stride=(1, 1))
def forward(self, x):
return self.unet(x)
#3a)------ --------- ---------
model_spine = SpineSegmentationModel2().to(device) #Kjører load_state_dict internt
model_spine.state_dict()
model_spine.load_state_dict(torch.load(pretrain_in))
I have uploaded fish.pth and the script for training, evaluating the model and 3 images if that helps here: https://github.com/ErlendAQK/Questions
I tried to alter just the last layer, but it must be altered on a “deeper” level: model_spine.unet.model[-1][0].conv = nn.ConvTranspose2d(32, 3, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), output_padding=(1, 1)) after model_spine.load_state_dict(…). I tried to make a model with 3 out_channels, but then I could not load the weights from the pretrained model.
Full code:
# -*- coding: utf-8 -*-
import torch
import monai
import torch.nn as nn
import torch.optim as optim
from monai.networks.nets import UNet
from monai.losses import DiceLoss
from monai.metrics import DiceMetric
from monai.data import Dataset, DataLoader
from monai.transforms import Compose, LoadImaged, EnsureChannelFirstd, ScaleIntensityd, RandCropByPosNegLabeld, RandRotate90d, SqueezeDimd
from monai.inferers import sliding_window_inference
from torch.utils.tensorboard import SummaryWriter
from monai.visualize import plot_2d_or_3d_image
import os
import sys
import logging
import numpy as np
import matplotlib.pyplot as plt
# Inputdata
datadir = r'C:UsersJohndatasetsimages'
segdir = r'C:UsersJohndatasetsimageslabels'
Nepochs = 50
pretrain_in ='fish.pth'
model_out = 'spine.pth'
def main(datadir, segdir, Nepochs, pretrain_in, model_out):
# ====================== STARTER ======================
monai.config.print_config()
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
# 1) --------- Leser inn bilder og labels ---------------------
objekter = [fil.removesuffix('.nii.gz') for fil in os.listdir(segdir) if fil.endswith('.nii.gz')]
bilder_inn = []
for obj in objekter:
path_img = os.path.join(datadir, obj + '.dcm.nii.gz')
path_seg = os.path.join(segdir, obj + '.nii.gz')
if os.path.exists(path_img) & os.path.exists(path_seg):
bilder_inn.append({'img': path_img, 'seg': path_seg})
else:
print(f"OPS! Fant ikke filene for '{obj}' - Hopper over.")
# Split dataset into training and validation
split_index = int(len(bilder_inn) * 0.8)
train_data_dicts = bilder_inn[:split_index]
val_data_dicts = bilder_inn[split_index:]
# 2) --------- Transforms ---------------------
train_transforms = Compose([
LoadImaged(keys=["img", "seg"], reader='PydicomReader'),
EnsureChannelFirstd(keys=["img", "seg"]),
SqueezeDimd(keys=["img", "seg"], dim=-1),
ScaleIntensityd(keys=["img"]),
RandCropByPosNegLabeld(keys=["img", "seg"], label_key="seg", spatial_size=[256, 256], pos=1, neg=1, num_samples=2),
RandRotate90d(keys=["img", "seg"], prob=0.5, spatial_axes=[0, 1]),
])
val_transforms = Compose([
LoadImaged(keys=["img", "seg"], reader='PydicomReader'),
EnsureChannelFirstd(keys=["img", "seg"]),
SqueezeDimd(keys=["img", "seg"], dim=-1),
ScaleIntensityd(keys=["img"]),
])
train_ds = Dataset(data=train_data_dicts, transform=train_transforms)
val_ds = Dataset(data=val_data_dicts, transform=val_transforms)
train_loader = DataLoader(train_ds, batch_size=1, shuffle=True, num_workers=4)
val_loader = DataLoader(val_ds, batch_size=1, num_workers=4)
# 3) --------- Model for spine segmentation ---------------
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
class SpineSegmentationModel2(nn.Module):
def __init__(self):
super(SpineSegmentationModel2, self).__init__()
self.unet = UNet(
spatial_dims=2,
in_channels=1, # 1-channel (grayscale input)
out_channels=2, # 2 classes (background, spine)
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
)
# Load pretrained weights
self.unet.load_state_dict(torch.load(pretrain_path), strict=False)
# Adjust to 3 output layers
self.unet.out_conv = nn.Conv2d(256, 3, kernel_size=(1, 1), stride=(1, 1))
def forward(self, x):
return self.unet(x)
#3a)------ --------- ---------
model_spine = SpineSegmentationModel2().to(device) #Kjører load_state_dict internt
model_spine.state_dict()
model_spine.load_state_dict(torch.load(pretrain_in))
# 4) --------- Loss function, optimizer og metrics ---------------
loss_function = nn.CrossEntropyLoss()
optimizer = optim.Adam(model_spine.parameters(), lr=1e-4)
dice_metric = DiceMetric(include_background=False, reduction="mean")
# 5) --------- Training loop ---------------
val_interval = 2
writer = SummaryWriter('runs_cascade') #(log_dir='runs/experiment_name')
for epoch in range(Nepochs):
print(f"Epoch {epoch+1}/{Nepochs}")
model_spine.train() # Set model to training mode
epoch_loss = 0
step = 0
for batch_data in train_loader:
step += 1
inputs = batch_data["img"].to(device)
labels = batch_data["seg"].to(device).long() #Konverterer til heltall
optimizer.zero_grad()
outputs = model_spine(inputs)
# Compute loss
loss = loss_function(outputs, labels.squeeze(1))
loss.backward()
optimizer.step()
epoch_loss += loss.item()
epoch_loss /= step
print(f"Training loss for epoch {epoch+1}: {epoch_loss}")
writer.add_scalar('Loss/train', epoch_loss / len(train_loader), epoch)
# Validering
if (epoch + 1) % val_interval == 0:
model_spine.eval()
with torch.no_grad():
val_dice = 0
val_steps = 0
for val_data in val_loader:
val_steps += 1
val_inputs = val_data["img"].to(device)
val_labels = val_data["seg"].to(device)
val_outputs = sliding_window_inference(val_inputs, roi_size=(256, 256), sw_batch_size=1, predictor=model_spine)
dice_metric(y_pred=val_outputs, y=val_labels)
mean_dice = dice_metric.aggregate().item()
dice_metric.reset()
print(f"Validation Dice score for epoch {epoch+1}: {mean_dice}")
# Lagre den trente modellen
torch.save(model_spine.state_dict(), model_out)
if __name__ == '__main__':
main(datadir, segdir, Nepochs, pretrain_in, model_out)