Help Implementing Checkpoint Loader?

[jmenville11]

Hi,

I'm trying to write my code so that my model saves every 50 epochs. Then, if the code is interrupted for whatever reason, I can load the last checkpoint and continue training from there. However, I'm having some issues with tensor agreement. Can you advise if I'm implementing CheckpointLoader wrong?

Jesse

# Setup imports
import os
import shutil
import tempfile
import glob
import nibabel as nib
import time
from pathlib import Path

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

import numpy as np
from numpy import array
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader

from monai.apps import DecathlonDataset
from monai.config import print_config
from monai.data import CacheDataset, DataLoader, decollate_batch
from monai.handlers import CheckpointLoader, ClassificationSaver, StatsHandler
from monai.handlers.utils import from_engine
from monai.losses import DiceLoss
from monai.inferers import sliding_window_inference
from monai.metrics import DiceMetric
from monai.networks.nets import SegResNet
from monai.transforms import (
    Activations,
    Activationsd,
    AsDiscrete,
    AsDiscreted,
    Compose,
    Invertd,
    LoadImaged,
    MapTransform,
    NormalizeIntensityd,
    Orientationd,
    RandFlipd,
    RandScaleIntensityd,
    RandShiftIntensityd,
    RandSpatialCropd,
    ResizeWithPadOrCropd,
    Spacingd,
    EnsureChannelFirstd,
    EnsureTyped,
    EnsureType,
)
from monai.utils import set_determinism

import torch
print_config()


# Setup data directory
data_dir = Path("/users/jmenvill/ICH/Model_Directory_256")
model_dir = Path("/gpfs/home/jmenvill/ICH/Model_Outputs_Jan_2022")
os.chdir(data_dir)


# Establish file paths
train_images = sorted(glob.glob(os.path.join(data_dir, "train", "Images", "*_FLAIR.nii")))
train_labels = sorted(glob.glob(os.path.join(data_dir, "train", "Labels", "*_label.nii")))
val_images = sorted(glob.glob(os.path.join(data_dir, "val", "Images", "*_FLAIR.nii")))
val_labels = sorted(glob.glob(os.path.join(data_dir, "val", "Labels", "*_label.nii")))

print('Training images found: ',len(train_images))
print('Training segmentations found: ',len(train_labels))
print('Validation images found: ',len(val_images))
print('Validation segmentations found: ',len(val_labels))


# Zip labels and images together into a dictionary for both training and validation sets
train_files = [
    {"image": image_name, "label": label_name}
    for image_name, label_name in zip(train_images, train_labels)
]
val_files = [
    {"image": image_name, "label": label_name}
    for image_name, label_name in zip(val_images, val_labels)
]


# Set deterministic training for reproducibility
set_determinism(seed=0)


# Convert brain tumor labels -> PHE and ICH labels
class ConvertToMultiChannelBasedOnBratsClassesd(MapTransform):
    """
    Convert labels to multi channels:
    label 1 is the hematoma (HE)
    label 2 is the edema (ED)
    The possible classes are ED or HE

    """
    def __call__(self, data):
        d = dict(data)
        for key in self.keys:
            result = []
            # label 1 is HE (hematoma)
            result.append(d[key] == 1)
            # label 2 is ED (edema)
            result.append(d[key] == 2)
            d[key] = np.stack(result, axis=0).astype(np.float32)
        return d 


# Setup training and validation transforms
train_transform = Compose(
    [
        LoadImaged(keys=["image", "label"]),
        EnsureChannelFirstd(keys="image"),
        ConvertToMultiChannelBasedOnBratsClassesd(keys="label"),
         ResizeWithPadOrCropd(keys=["image", "label"],
             spatial_size=(256,256,256),
             mode="constant"),
        Spacingd(
            keys=["image", "label"],
            pixdim=(1.0, 1.0, 1.0),
            mode=("bilinear", "nearest"),
        ),
        Orientationd(keys=["image", "label"], axcodes="RAS"),
        RandSpatialCropd(keys=["image", "label"], roi_size=[224, 224, 144], random_size=False),
        RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=0),
        RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=1),
        RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=2),
        NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
        RandScaleIntensityd(keys="image", factors=0.1, prob=1.0),
        RandShiftIntensityd(keys="image", offsets=0.1, prob=1.0),
        EnsureTyped(keys=["image", "label"]),
    ]
)
val_transform = Compose(
    [
        LoadImaged(keys=["image", "label"]),
        EnsureChannelFirstd(keys="image"),
        ConvertToMultiChannelBasedOnBratsClassesd(keys="label"),
         ResizeWithPadOrCropd(keys=["image", "label"],
             spatial_size=(256,256,256),
             mode="constant"),
        Spacingd(
            keys=["image", "label"],
            pixdim=(1.0, 1.0, 1.0),
            mode=("bilinear", "nearest"),
        ),
        Orientationd(keys=["image", "label"], axcodes="RAS"),
        NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
        EnsureTyped(keys=["image", "label"]),
    ]
)


# Establish training and validation datasets
train_ds = CacheDataset(
    data=train_files,
    transform=train_transform,
    cache_rate=0,
    num_workers=2)

train_loader = DataLoader(
    train_ds,
    batch_size=1,
    shuffle=True,
    num_workers=2) 

val_ds = CacheDataset(
    data=val_files,
    transform=val_transform,
    cache_rate=0,
    num_workers=2)

val_loader = DataLoader(
    val_ds,
    batch_size=1,
    shuffle=False,
    num_workers=0) 

print(f"train image shape: {train_ds[12]['image'].shape}")
print(f"train label shape: {train_ds[12]['label'].shape}")
print(f"val image shape: {val_ds[4]['image'].shape}")
print(f"val label shape: {val_ds[4]['label'].shape}")


# Pick one image from dataset to visualize and check label channels
# print(f"image shape: {val_ds[2]['image'].shape}")
# plt.figure("image", (24, 6))
# for i in range(1):
#    plt.subplot(1, 3, i + 1)
#    plt.title(f"image channel {i}")
#    plt.imshow(val_ds[2]["image"][i, :, :, 40].detach().cpu(), cmap="gray")
# plt.show()

# Load corresponding label channels (2) for this image
# print(f"label shape: {val_ds[2]['label'].shape}")
# plt.figure("label", (24, 6))
# for i in range(2):
#    plt.subplot(1, 2, i + 1)
#    plt.title(f"label channel {i}")
#    plt.imshow(val_ds[2]["label"][i, :, :, 40].detach().cpu())
# plt.show()


# Set epochs
max_epochs = 3
val_interval = 1
VAL_AMP = True


# Load checkpoint and define device
load_checkpoint = True
device = torch.device("cuda:0")


# Define model
model = SegResNet(
        blocks_down=[1, 2, 2, 4],
        blocks_up=[1, 1, 1],
        init_filters=16,
        in_channels=1,
        out_channels=2,
        dropout_prob=0.2,
    ).to(device)


if load_checkpoint:
    model.load_state_dict(torch.load("/gpfs/home/jmenvill/ICH/Model_Outputs_Jan_2022/model_checkpoint_saved.pth"))
    model.eval()
    model.to(device)
else:
    model = SegResNet(
        blocks_down=[1, 2, 2, 4],
        blocks_up=[1, 1, 1],
        init_filters=16,
        in_channels=1,
        out_channels=2,
        dropout_prob=0.2,
    ).to(device)

loss_function = DiceLoss(smooth_nr=0, smooth_dr=1e-5, squared_pred=True, to_onehot_y=False, sigmoid=True)
optimizer = torch.optim.Adam(model.parameters(), 1e-4, weight_decay=1e-5)
lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=max_epochs)

dice_metric = DiceMetric(include_background=True, reduction="mean")
dice_metric_batch = DiceMetric(include_background=True, reduction="mean_batch")

post_trans = Compose(
    [EnsureType(), Activations(sigmoid=True), AsDiscrete(threshold_values=True)]
)


# Define inference method
def inference(input):
    
    def _compute(input):
        return sliding_window_inference(
            inputs=input,
            roi_size=(240, 240, 160),
            sw_batch_size=1,
            predictor=model,
            overlap=0.5,
        )

    if VAL_AMP:
        with torch.cuda.amp.autocast():
            return _compute(input)
    else:
        return _compute(input)

scaler = torch.cuda.amp.GradScaler()
torch.backends.cudnn.benchmark = True


# Execute training
best_metric = -1
best_metric_epoch = -1
best_metrics_epochs_and_time = [[], [], []]
epoch_loss_values = []
metric_values = []
metric_values_he = []
metric_values_ed = []

total_start = time.time()
for epoch in range(max_epochs):
    if epoch %50 == 0: 
        torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss_function,
            }, os.path.join(model_dir, "model_checkpoint_saved.pth"))
	
    epoch_start = time.time()
    print("-" * 10)
    print(f"epoch {epoch + 1}/{max_epochs}")
    model.train()
    epoch_loss = 0
    step = 0
    for batch_data in train_loader:
        step_start = time.time()
        step += 1
        inputs, labels = (
            batch_data["image"].to(device),
            batch_data["label"].to(device),
        )
        optimizer.zero_grad()
        with torch.cuda.amp.autocast():
            outputs = model(inputs)
            loss = loss_function(outputs, labels)
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
        epoch_loss += loss.item()
        print(
            f"{step}/{len(train_ds) // train_loader.batch_size}"
            f", train_loss: {loss.item():.4f}"
            f", step time: {(time.time() - step_start):.4f}"
        )
    lr_scheduler.step()
    epoch_loss /= step
    epoch_loss_values.append(epoch_loss)
    print(f"epoch {epoch + 1} average loss: {epoch_loss:.4f}")

    if (epoch + 1) % val_interval == 0:
        model.eval()
        with torch.no_grad():

            for val_data in val_loader:
                val_inputs, val_labels = (
                    val_data["image"].to(device),
                    val_data["label"].to(device),
                )
                val_outputs = inference(val_inputs)
                val_outputs = [post_trans(i) for i in decollate_batch(val_outputs)]
                dice_metric(y_pred=val_outputs, y=val_labels)
                dice_metric_batch(y_pred=val_outputs, y=val_labels)

            metric = dice_metric.aggregate().item()
            metric_values.append(metric)
            metric_batch = dice_metric_batch.aggregate()
            
            metric_he = metric_batch[0].item()
            metric_values_he.append(metric_he)
            
            metric_ed = metric_batch[1].item()
            metric_values_ed.append(metric_ed)

            dice_metric.reset()
            dice_metric_batch.reset()

            if metric > best_metric:
                best_metric = metric
                best_metric_epoch = epoch + 1
                best_metrics_epochs_and_time[0].append(best_metric)
                best_metrics_epochs_and_time[1].append(best_metric_epoch)
                best_metrics_epochs_and_time[2].append(time.time() - total_start)
                torch.save(
                    model.state_dict(),
                    os.path.join(model_dir, "best_model.pth"),
                )
                print("saved new best metric model")
            print(
                f"current epoch: {epoch + 1} current mean dice: {metric:.4f}"
                f" hematoma: {metric_he:.4f} edema: {metric_ed:.4f}"
                f"\nbest mean dice: {best_metric:.4f}"
                f" at epoch: {best_metric_epoch}"
            )
    print(f"time consuming of epoch {epoch + 1} is: {(time.time() - epoch_start):.4f}")
total_time = time.time() - total_start

print(f"train completed, best_metric: {best_metric:.4f} at epoch: {best_metric_epoch}, total time: {total_time}.")



# Save training metrics in plot form
plt.figure("train", (12, 6))
plt.subplot(1, 2, 1)
plt.title("Epoch Average Loss")
x = [i + 1 for i in range(len(epoch_loss_values))]
y = epoch_loss_values
plt.xlabel("epoch")
plt.plot(x, y, color="purple")

plt.subplot(1, 2, 2)
plt.title("Val Mean Dice")
x = [val_interval * (i + 1) for i in range(len(metric_values))]
y = metric_values
plt.xlabel("epoch")
plt.plot(x, y, color="green")
# plt.show()
plt.savefig('/users/jmenvill/ICH/Model_Outputs_Jan_2022/Epoch_Average_Loss_&_Val_Mean_Dice.png', bbox_inches='tight')

plt.figure("train", (18, 6))
plt.subplot(1, 3, 1)
plt.title("Hematoma Mean Dice")
x = [val_interval * (i + 1) for i in range(len(metric_values_he))]
y = metric_values_he
plt.xlabel("epoch")
plt.plot(x, y, color="red")

plt.subplot(1, 3, 2)
plt.title("Val Mean Dice Edema")
x = [val_interval * (i + 1) for i in range(len(metric_values_ed))]
y = metric_values_ed
plt.xlabel("epoch")
plt.plot(x, y, color="blue")
# plt.show()
plt.savefig('/users/jmenvill/ICH/Model_Outputs_Jan_2022/Val_Mean_Dice_Hematoma_&_Edema.png', bbox_inches='tight')




# Check best model against validation set
model.load_state_dict(
    torch.load(os.path.join(model_dir, "best_model.pth"))
)
model.eval()
with torch.no_grad():
    
# Select one image to evaluate and visualize the model output
    val_input = val_ds[6]["image"].unsqueeze(0).to(device)
    roi_size = (128, 128, 64)
    sw_batch_size = 4
    val_output = inference(val_input)
    val_output = post_trans(val_output[6])
    plt.figure("image", (24, 6))
    for i in range(1):
        plt.subplot(1, 1, i + 1)
        plt.title(f"image channel {i}")
        plt.imshow(val_ds[6]["image"][i, :, :, 100].detach().cpu(), cmap="gray")
    # plt.show()
    plt.savefig('users/jmenvill/ICH/Model_Outputs_Jan_2022/Model_Performance_Check_Image.png', bbox_inches='tight')
    
# Visualize the 3 label channels corresponding to this image (ground truth)
    plt.figure("label", (24, 6))
    for i in range(2):
        plt.subplot(1, 2, i + 1)
        plt.title(f"label channel {i}")
        plt.imshow(val_ds[6]["label"][i, :, :, 100].detach().cpu())
    # plt.show()
    plt.savefig('users/jmenvill/ICH/Model_Outputs_Jan_2022/Model_Performance_Check_Labels.png', bbox_inches='tight')
    
# Visualize the model outputs corresponding to this image (model output predictions)
    plt.figure("output", (18, 6))
    for i in range(2):
        plt.subplot(1, 2, i + 1)
        plt.title(f"output channel {i}")
        plt.imshow(val_output[i, :, :, 100].detach().cpu())
    # plt.show()
    plt.savefig('users/jmenvill/ICH/Model_Outputs_Jan_2022/Model_Performance_Check_Outputs.png', bbox_inches='tight')](url)

Error:

[Nic-Ma]

Hi @jmenville11 ,

Seems you saved a dict of data to ckpt, please only load the model_state_dict to apply to your model.

Thanks.

[dannyhow12]

Hi and good day @jmenville11 ,

Thank you for the wonderful question! I was able to reproduce it and suspected it could be due to the arrangement of model.to(device). Below is the code that I have changed from your original code and the error that prompts.

model = SegResNet(
        blocks_down=[1, 2, 2, 4],
        blocks_up=[1, 1, 1],
        init_filters=16,
        in_channels=1,
        out_channels=2,
        dropout_prob=0.2,
    )

if load_checkpoint:
    model.load_state_dict(torch.load("/gpfs/home/jmenvill/ICH/Model_Outputs_Jan_2022/model_checkpoint_saved.pth"))
    model.eval()
    model.to(device)
else:
    model.to(device)

Also, I believe @Nic-Ma was referencing to this. It can be referred from this link https://pytorch.org/tutorials/recipes/recipes/saving_and_loading_a_general_checkpoint.html . The code I have attached down here is from the link, where you may refer to for better clarity.

checkpoint = torch.load("/.../model_checkpoint_saved.pth")
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epochs = checkpoint['epoch']
loss = checkpoint['loss']
# model.eval()
model.train()
model.to(device)

This should now solve the problem. Credits to @Nic-Ma for pointing it out. I am leaving this here in case anyone encounters similar issue before.

Regards,
Danny