GPU usage become lower after a few epochs

[fanqiehc]

I'm training an unet model to do 3D segmentation job.
There are some performance issues that CPU usage becomes higher and GPU usage becomes lower after a few epochs training.

The usage metrics (recorded by Neptune) are shown as blow:

My code are shown as below:

set_determinism(seed=111)

import os

from monai.losses import DiceCELoss
from monai.inferers import sliding_window_inference
from monai.transforms import (
    AsDiscrete,
    AddChanneld,
    Compose,
    LoadImaged,
    RandFlipd,
    RandCropByPosNegLabeld,
    RandShiftIntensityd,
    ScaleIntensityRanged,
    RandRotate90d,
    ToTensord,
    SpatialPadd,
    CenterSpatialCropd,
)

from monai.config import print_config
from monai.metrics import DiceMetric
from monai.networks.nets import UNETR, UNet

from monai.data import (
    CacheDataset,
    load_decathlon_datalist,
    decollate_batch,
    list_data_collate,
)

import torch
import pytorch_lightning
from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint

from pytorch_lightning.loggers import NeptuneLogger

neptune_logger = NeptuneLogger(
    project="",
    api_token="",
)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
torch.backends.cudnn.benchmark = True
print_config()

directory = os.environ.get("MONAI_DATA_DIRECTORY")
root_dir = "./models" #tempfile.mkdtemp() if directory is None else directory
print(root_dir)


class Net(pytorch_lightning.LightningModule):
    def __init__(self):
        super().__init__()

        self.input_size = (144, 144, 144)

        NumClasses = 8

        self._model = UNet(
            spatial_dims=3,
            in_channels=1,
            out_channels=NumClasses,
            channels=(16, 32, 64, 128),
            strides=(2, 2, 2),
            num_res_units=2,
        ).to(device)

        self.loss_function = DiceCELoss(to_onehot_y=True, softmax=True)
        self.post_pred = AsDiscrete(argmax=True, to_onehot=NumClasses)
        self.post_label = AsDiscrete(to_onehot=NumClasses)
        self.dice_metric = DiceMetric(
            include_background=True, reduction="mean", get_not_nans=False
        )
        self.best_val_dice = 0
        self.best_val_epoch = 0
        self.max_epochs = 100
        self.check_val = 1
        self.warmup_epochs = 1
        self.metric_values = []
        self.epoch_loss_values = []

    def forward(self, x):
        return self._model(x)

    def prepare_data(self):
        # prepare data
        datasets = "./dataset_fold1.json"
        datalist = load_decathlon_datalist(datasets, True, "training")
        val_files = load_decathlon_datalist(datasets, True, "validation")

        train_transforms = Compose(
            [
                LoadImaged(keys=["image", "label"]),
                AddChanneld(keys=["image", "label"]),
                ScaleIntensityRanged(
                    keys=["image"],
                    a_min=-200,
                    a_max=1000,
                    b_min=0.0,
                    b_max=1.0,
                    clip=True,
                ),
                SpatialPadd(
                    keys=["image", "label"],
                    spatial_size=self.input_size,
                ),
                RandCropByPosNegLabeld(
                    keys=["image", "label"],
                    label_key="label",
                    pos=1,
                    neg=1,
                    spatial_size=self.input_size,
                    num_samples=4,
                ),
                RandFlipd(
                    keys=["image", "label"],
                    spatial_axis=[0],
                    prob=0.10,
                ),
                RandFlipd(
                    keys=["image", "label"],
                    spatial_axis=[1],
                    prob=0.10,
                ),
                RandFlipd(
                    keys=["image", "label"],
                    spatial_axis=[2],
                    prob=0.10,
                ),
                RandRotate90d(
                    keys=["image", "label"],
                    prob=0.10,
                    max_k=3,
                ),
                RandShiftIntensityd(
                    keys=["image"],
                    offsets=0.10,
                    prob=0.50,
                ),
                ToTensord(keys=["image", "label"]),
            ]
        )
        val_transforms = Compose(
            [
                LoadImaged(keys=["image", "label"]),
                AddChanneld(keys=["image", "label"]),
                SpatialPadd(
                    keys=["image", "label"],
                    spatial_size=self.input_size,
                ),
                CenterSpatialCropd(
                    keys=["image", "label"],
                    roi_size=self.input_size,
                ),
                ScaleIntensityRanged(
                    keys=["image"],
                    a_min=-200,
                    a_max=1000,
                    b_min=0.0,
                    b_max=1.0,
                    clip=True,
                ),
                ToTensord(keys=["image", "label"]),
            ]
        )

        self.train_ds = CacheDataset(
            data=datalist,
            transform=train_transforms,
            cache_num=24,
            cache_rate=1.0,
            num_workers=12,
        )
        self.val_ds = CacheDataset(
            data=val_files,
            transform=val_transforms,
            cache_num=6,
            cache_rate=1.0,
            num_workers=8,
        )

    def train_dataloader(self):
        train_loader = torch.utils.data.DataLoader(
            self.train_ds,
            batch_size=4,
            shuffle=True,
            num_workers=16,
            pin_memory=True,
            collate_fn=list_data_collate,
        )
        return train_loader

    def val_dataloader(self):
        val_loader = torch.utils.data.DataLoader(
            self.val_ds, batch_size=4, shuffle=False, num_workers=16, pin_memory=True
        )
        return val_loader

    def configure_optimizers(self):
        optimizer = torch.optim.AdamW(
            self._model.parameters(), lr=1e-4, weight_decay=1e-5
        )
        return optimizer

    def training_step(self, batch, batch_idx):
        images, labels = (batch["image"].cuda(), batch["label"].cuda())
        output = self.forward(images)
        loss = self.loss_function(output, labels)
        tensorboard_logs = {"train_loss": loss.item()}
        self.log('train_loss', loss, prog_bar=False, on_epoch=True, on_step=True)
        self.log("batch_loss", loss, prog_bar=False, on_epoch=True, on_step=True)
        return {"loss": loss, "log": tensorboard_logs}

    def training_epoch_end(self, outputs):
        avg_loss = torch.stack([x["loss"] for x in outputs]).mean()
        self.log('avg_loss', avg_loss, on_epoch=True, on_step=False)
        self.epoch_loss_values.append(avg_loss.detach().cpu().numpy())

    def validation_step(self, batch, batch_idx):
        images, labels = batch["image"], batch["label"]
        roi_size = self.input_size
        sw_batch_size = 4
        outputs = sliding_window_inference(
            images, roi_size, sw_batch_size, self.forward
        )
        loss = self.loss_function(outputs, labels)
        outputs = [self.post_pred(i) for i in decollate_batch(outputs)]
        labels = [self.post_label(i) for i in decollate_batch(labels)]
        dice = self.dice_metric(y_pred=outputs, y=labels)

        self.log('val_loss', loss, prog_bar=False, on_step=True, on_epoch=True)
        return {"val_loss": loss, "val_number": len(outputs)}

    def validation_epoch_end(self, outputs):
        val_loss, num_items = 0, 0
        for output in outputs:
            val_loss += output["val_loss"].sum().item()
            num_items += output["val_number"]
        mean_val_dice = self.dice_metric.aggregate().item()
        self.dice_metric.reset()
        mean_val_loss = torch.tensor(val_loss / num_items)
        tensorboard_logs = {
            "val_dice": mean_val_dice,
            "val_loss": mean_val_loss,
        }
        if mean_val_dice > self.best_val_dice:
            self.best_val_dice = mean_val_dice
            self.best_val_epoch = self.current_epoch
        print(
            f"\ncurrent epoch: {self.current_epoch} "
            f"current mean dice: {mean_val_dice:.4f}"
            f"\nbest mean dice: {self.best_val_dice:.4f} "
            f"at epoch: {self.best_val_epoch}"
        )
        self.log('val_avg_loss', mean_val_loss, prog_bar=False, on_step=False, on_epoch=True)
        self.log('val_avg_dice', mean_val_dice, prog_bar=False, on_step=False, on_epoch=True)
        self.metric_values.append(mean_val_dice)
        return {"val_loss": mean_val_loss, "log": tensorboard_logs}

# initialise the LightningModule
net = Net()

# set up checkpoints
checkpoint_callback = ModelCheckpoint(
    monitor='val_loss',
    auto_insert_metric_name=False,
    dirpath=root_dir,
    filename='unet-epoch{epoch:02d}-val_loss{val_loss:.2f}',
    save_top_k=50, 
    every_n_epochs=1,
    verbose=True,
)

# initialise Lightning's trainer.
trainer = pytorch_lightning.Trainer(
    gpus=[0] if torch.cuda.is_available() else [],
    max_epochs=net.max_epochs,
    check_val_every_n_epoch=net.check_val,
    callbacks=[checkpoint_callback],
    default_root_dir=root_dir,
    logger=neptune_logger,
)

# train
trainer.fit(net)

neptune_logger.experiment.stop()

[Nic-Ma]

Hi @fanqiehc ,

I checked your charts, seems your CPU usage and memory usage started to change at epoch 500, was there anything special at that time? Or do you always get the same training curve every time?

Thanks in advance.

[fanqiehc]

I haven't found any thing happened at that time. I tried serval times and got different curves.

Here is a normal one

[Nic-Ma]

Thanks for the quick testing.
I feel it seems not related to the training program, maybe your environment is not very stable?
For example, may have some other programs running or windows/linux GUI opened?

Thanks.

[fanqiehc]

The system is almost clean, ubuntu 20.04, training inside docker container, no other program running.
But maybe you are right, I've seen many unstable hardwares.
I will do more test and compare it with cloud environment.

Thanks for your response.