Loss cannot converge in binary classification using DenseNet121.

[Minxiangliu]

I try to use DenseNet121 model for binary classification problem in 3D, and I refer to spleen_segmentation_3d_lightning.ipynb and densenet_training_array.ipynb code in tutorials. Everything is fine, but the result is confusing, whether my Are there any errors in the code?

Thanks in advance!

class Net(pytorch_lightning.LightningModule):
    def __init__(self, tr_path, te_path, train_pad_shape, train_crop_shape, test_pad_shape):
        super().__init__()

        self.tr_path = tr_path
        self.te_path = te_path

        self.train_pad_shape = train_pad_shape
        self.train_crop_shape = train_crop_shape
        self.test_pad_shape = test_pad_shape 

        self.model = DenseNet121(
            spatial_dims=3, 
            in_channels=1, 
            out_channels=1
        ).to(DEVICE)
   
        self.activation_sigmoid = torch.nn.Sigmoid()   
        self.loss_binary_cross_Entropy = torch.nn.BCELoss()
        
        self.best_val_acc = 0
        self.best_val_epoch = 0
    
        self.sum_train_loss = 0  
        self.sum_val_loss = 0
        self.sum_train_correct = 0
        self.sum_val_correct = 0
                
    def forward(self, x):
        return self.activation_sigmoid(self.model(x))
    
    def prepare_data(self):
        tr_ = np.loadtxt(self.tr_path, delimiter=',', skiprows=1, usecols=[1,2,3,4], dtype=np.str_, encoding="utf-8")
        te_ = np.loadtxt(self.te_path, delimiter=',', skiprows=1, usecols=[1,2,3,4], dtype=np.str_, encoding="utf-8")

        train_files = [
            {"image": image_name, 'label':label_name}
            for image_name, label_name in zip(tr_[:,0].reshape(-1,1), 
                                              tr_[:,1].reshape(-1,1).astype(np.float32))
        ]
        
        test_files = [
            {"image": image_name, 'label':label_name}
            for image_name, label_name in zip(te_[:,0].reshape(-1,1), 
                                              te_[:,1].reshape(-1,1).astype(np.float32))
        ]

        # set deterministic training for reproducibility
        set_determinism(seed=0)
        
        # define the data transforms
        train_transforms = [LoadImaged(keys=['image']),
            EnsureChannelFirstd(keys=['image']),
            Spacingd(keys=['image'], pixdim=(4., 4., 1.), mode=('bilinear')),
            Orientationd(keys=['image'], axcodes='RAS'),
            Rotate90d(keys=['image'], spatial_axes=(0,1)),
            ScaleIntensityd(keys=['image']),
            CropForegroundd(keys=['image'], source_key='image', select_fn=lambda x:x>0.05, margin=0, return_coords=False),
            SpatialPadd(keys=['image'], spatial_size=self.train_pad_shape),
            EnsureTyped(keys=['image', 'label']),
            ToDeviced(keys=['image', 'label'], device='cuda:0')]

        train_transforms.extend([
            RandSpatialCropSamplesd(keys=['image'], num_samples=8, roi_size=self.train_crop_shape, random_size=False),
            RandFlipd(keys=['image'],prob=0.5),
            RandRotate90d(keys=['image'],prob=0.5),
            RandScaleIntensityd(keys=['image'],prob=0.5,factors=(5,10)),
            RandShiftIntensityd(keys=['image'],prob=0.5,offsets=(10, 20)),
            RandGibbsNoised(keys=['image'],prob=0.5,alpha=(0.2, 0.8)),
            RandKSpaceSpikeNoised(keys=['image'],prob=0.5,intensity_range=(8,10)),
            RandCoarseDropoutd(keys=['image'],prob=0.5,spatial_size=10, holes=100, fill_value=0),
            RandCoarseShuffled(keys=['image'],prob=0.5,spatial_size=10, holes=100)
        ])
        
        test_transforms = Compose([LoadImaged(keys=['image']),
            EnsureChannelFirstd(keys=['image']),
            Spacingd(keys=['image'], pixdim=(4., 4., 1.), mode=('bilinear')),
            Orientationd(keys=['image'], axcodes='RAS'),
            Rotate90d(keys=['image'], spatial_axes=(0,1)),
            ScaleIntensityd(keys=['image']),
            CropForegroundd(keys=['image'], source_key='image', select_fn=lambda x:x>0.05, margin=0, return_coords=False),
            SpatialPadd(keys=['image'], spatial_size=self.test_pad_shape),
            EnsureTyped(keys=['image', 'label']),
            ToDeviced(keys=['image', 'label'], device='cuda:0')])
        
        self.train_ds = CacheDataset(
            data=train_files, 
            transform=Compose(train_transforms),
            cache_rate=1.0,
            copy_cache=False
        )

        self.test_ds = CacheDataset(
            data=test_files, 
            transform=test_transforms,
            cache_rate=1.0, 
            copy_cache=False
        )
        
    def train_dataloader(self):
        train_loader = ThreadDataLoader(self.train_ds, 
                                        num_workers=0, 
                                        batch_size=batch_size, 
                                        shuffle=True)
        return train_loader

    def val_dataloader(self):
        val_loader = ThreadDataLoader(self.test_ds, 
                                      num_workers=0, 
                                      batch_size=3,
                                      shuffle=False)
        return val_loader

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.model.parameters(), 1e-4)
        return optimizer
    
    def training_step(self, batch, batch_idx):
        images, labels = batch["image"].to(DEVICE), batch["label"].to(DEVICE)
        output = self.forward(images)
        loss = self.loss_binary_cross_Entropy(output, labels)
        value = torch.eq(torch.gt(output, 0.5).to(torch.int32), labels)
        self.sum_train_correct += value.sum().item()
        return {'loss':loss, 'train_number': len(labels)}

    def training_epoch_end(self, training_step_outputs):
        num_items = 0
        for N, output in enumerate(training_step_outputs):
            self.sum_train_loss += output['loss']
            num_items += output['train_number']
        
        mean_train_loss = self.sum_train_loss/(N+1)
        mean_val_acc = self.sum_train_correct/num_items

        self.sum_train_loss = 0
        self.sum_train_correct = 0

        print("current epoch: %d"%self.current_epoch)
        print("current mean train loss: %.4f"%mean_train_loss)
        print("current mean train acc: %.4f"%mean_val_acc)
        tensorboard_logs = {'train_loss': mean_train_loss,
                            'train_acc':mean_val_acc}
        self.log("log", tensorboard_logs)

    def validation_step(self, batch, batch_idx):       
        images, labels = batch["image"].to(DEVICE), batch["label"].to(DEVICE)
        output = self.forward(images)
        val_loss = self.loss_binary_cross_Entropy(output, labels)
        value = torch.eq(torch.gt(output, 0.5).to(torch.int32), labels)
        self.sum_val_correct += value.sum().item()
        return {'val_loss':val_loss, 'val_number': len(labels)}
    
    def validation_epoch_end(self, validation_step_outputs):
        num_items = 0
        for N,output in enumerate(validation_step_outputs):
            self.sum_val_loss += output["val_loss"]
            num_items += output['val_number']
        
        mean_val_loss = self.sum_val_loss/(N+1)
        mean_val_acc = self.sum_val_correct/num_items
        
        self.sum_val_loss = 0
        self.sum_val_correct = 0

        tensorboard_logs = {
            "val_loss": mean_val_loss,
            "val_acc": mean_val_acc
            }
        self.log("log", tensorboard_logs)
        self.log("val_acc", mean_val_acc)

        if self.current_epoch > 5:
            if mean_val_acc > self.best_val_acc:
                self.best_val_acc = mean_val_acc
                self.best_val_epoch = self.current_epoch 

        print('*'*10,'validation','*'*10)
        print("current epoch: %d"%self.current_epoch)
        print("current mean val loss: %.4f"%mean_val_loss)
        print("current mean acc: %.4f"%mean_val_acc)
        print("best mean acc: %.4f at epoch: %d\n"%(self.best_val_acc, self.best_val_epoch))
        
if __name__ == '__main__':
    log_dir = os.path.join("....")
    tb_logger = TensorBoardLogger(
        save_dir=log_dir
    )

    checkpoint_dir = os.path.join(tb_logger.save_dir, tb_logger.name, 'version_%d'%tb_logger.version,'checkpoints')
    cp = ModelCheckpoint(
        monitor='val_acc',
        mode='max',
        save_weights_only=True,
        save_top_k=5,
        dirpath=checkpoint_dir,
        filename='{epoch:04d}-{val_acc:.2f}',
        )
    
    batch_size = 16
    net = Net(tr_path='tr_roi_nii.csv', 
              te_path='te_roi_nii.csv',
              train_pad_shape=(96,96,32),
              train_crop_shape=(64,64,32),
              test_pad_shape=(64,64,32))

    trainer = pytorch_lightning.Trainer(
        enable_progress_bar=False,
        gpus=[0],
        max_epochs=2000,
        logger=tb_logger,
        enable_checkpointing=True,
        num_sanity_val_steps=1,
        log_every_n_steps=10,
        check_val_every_n_epoch=10,
        callbacks=[cp]
    )
    trainer.fit(net)

Reuslts:

[yiheng-wang-nv]

Hi @Minxiangliu ,
First of all, could you put Orientationd before Spacingd?
Then, since the training loss shakes and cannot converge, my suggestion is to remove all image augmentation operations, and check if the training loss can converge (overfit the training set)?
Could you please have a try and update the results here, and then we can do a further analysis if needed.

[Minxiangliu]

Hi @yiheng-wang-nv , thanks for your reply.
I have put Orientationd in front of Spacingd, and then removed the augmentation function, the final result is as shown in the figure:

train_transforms = [LoadImaged(keys=['image']),
            EnsureChannelFirstd(keys=['image']),
            Orientationd(keys=['image'], axcodes='RAS'),
            Spacingd(keys=['image'], pixdim=(4., 4., 1.), mode=('bilinear')),
            Rotate90d(keys=['image'], spatial_axes=(0,1)),
            ScaleIntensityd(keys=['image']),
            SpatialPadd(keys=['image'], spatial_size=(96,96,32)),
            CenterSpatialCropd(keys=['image'], roi_size=(64,64,32)),
            EnsureTyped(keys=['image', 'label']),
            ToDeviced(keys=['image', 'label'], device='cuda:0')]

By the way, the train data and test data number is 17 and 3. (patient)

[yiheng-wang-nv]

Hi @Minxiangliu ,

It seems the training set can be overfitted, thus you can try to gradually add random transforms and check the validation scores.
Be careful since the dataset is very small, and I suggest to use cross validation rather than train test split.

In addition, what are the original sizes of all samples? I guess this transform RandSpatialCropSamplesd may bring some issues.

[Minxiangliu]

I will try cross-validation.
The transform RandSpatialCropSamplesd, I have confirmed that the input image is what I want.

[Minxiangliu]

Hi, @yiheng-wang-nv,
I haven't performed cross-validation yet. I first tried to gradually increase the augmentation function. During the process, I found that as long as I added RandShiftIntensityd, the training would start to shake, so that it could not converge. Do you have any comments on this?

Could it be that my data is not suitable for using RandShiftIntensityd or my parameters are set wrong,
The parameters I set are here.

train_transforms = [LoadImaged(keys=['image']),
            EnsureChannelFirstd(keys=['image']),
            Orientationd(keys=['image'], axcodes='RAS'),
            Spacingd(keys=['image'], pixdim=(4., 4., 1.), mode=('bilinear')),
            Rotate90d(keys=['image'], spatial_axes=(0,1)),
            ScaleIntensityd(keys=['image']),
            CropForegroundd(keys=['image'], source_key='image', select_fn=lambda x:x>0.05, margin=0, return_coords=False),
            SpatialPadd(keys=['image'], spatial_size=self.train_pad_shape),
            EnsureTyped(keys=['image', 'label']),
            ToDeviced(keys=['image', 'label'], device='cuda:0')]

        train_transforms.extend([
            RandSpatialCropSamplesd(keys=['image'], num_samples=8, roi_size=self.train_crop_shape, random_size=False),
            RandFlipd(keys=['image'],prob=0.5),
            RandRotate90d(keys=['image'],prob=0.5),
            RandScaleIntensityd(keys=['image'],prob=0.5,factors=(5,10)),
            # RandShiftIntensityd(keys=['image'],prob=0.5,offsets=(10, 20)),   ←start not converging properly
        ])
        
        test_transforms = Compose([LoadImaged(keys=['image']),
            EnsureChannelFirstd(keys=['image']),
            Spacingd(keys=['image'], pixdim=(4., 4., 1.), mode=('bilinear')),
            Orientationd(keys=['image'], axcodes='RAS'),
            Rotate90d(keys=['image'], spatial_axes=(0,1)),
            ScaleIntensityd(keys=['image']),
            CropForegroundd(keys=['image'], source_key='image', select_fn=lambda x:x>0.05, margin=0, return_coords=False),
            SpatialPadd(keys=['image'], spatial_size=self.test_pad_shape),
            EnsureTyped(keys=['image', 'label']),
            ToDeviced(keys=['image', 'label'], device='cuda:0')])

[yiheng-wang-nv]

I see,
you used ScaleIntensityd(keys=['image'], which convert the values of the input into [0, 1], thus RandShiftIntensityd(keys=['image'],prob=0.5,offsets=(10, 20)) have 50% prob to change the values into > 10 (too large!)
In addition, I think factors for RandScaleIntensityd(keys=['image'],prob=0.5,factors=(5,10)) is also not meaningful.
Could you please check the doc-strings for these two transforms?

[Minxiangliu]

Hi @yiheng-wang-nv ,
Do you mean to check the factors of RandShiftIntensityd and RandScaleIntensityd in the documentation?
RandShiftIntensityd mainly uses the offsets parameter to change the value.
RandScaleIntensityd mainly uses the v = v * (1 + factor) to change the value.
So as you said, after already using ScaleIntensityd, it may be more appropriate to use RandScaleIntensityd, and I should set the parameters myself instead of referring to this.

Or do you have a better opinion, what parameters should be set to use RandScaleIntensityd for values in the range 0~1?

[yiheng-wang-nv]

Hi @Minxiangliu , for RandScaleIntensityd , you can try start from a small value, like factors=(-0.1, 0.1), if it is helpful, you can also try to increase the value into (-0.2, 0.2) to see if it is more helpful to avoid overfitting.