How to use tuner.scale_batch_size correctly?

[markvvw]

I tried to used tuner.scale_batch_size to find the optimal batch size in my work, however, the tuner.scale_batch_size returns a suitable number (obviously bigger) . How to used the funcion of pl correctly? thx

---

My code is the following:

class LitVnet(lit.LightningModule):
    def __init__(self, img_size, net, loss_fn, lr, batch_size, epoch):
        super().__init__()
        self.img_size = img_size
        self.net = net
        self.loss_fn = loss_fn
        self.lr = lr
        self.batch_size = batch_size
        self.epoch = epoch
        self.epoch_loss_values = []
        self.dice_values = []
        self.dice = DiceMetric(reduction="mean", include_background=True)
        self.post_trans = Compose(
            [EnsureType(), Activations(sigmoid=True), AsDiscrete(threshold=0.5)])

    def forward(self, x):
        return self.net(x)
    def train_dataloader(self):
        seg_train = SegDataset('./data/seg/train/',
                               train=True, size=self.img_size)
        return DataLoader(seg_train, batch_size=self.batch_size, shuffle=True)

    def val_dataloader(self):
        seg_val = SegDataset(
            './data/seg/val/', train=False, size=self.img_size)
        return DataLoader(seg_val, batch_size=self.batch_size)

    def configure_optimizers(self):
        optimizer = optim.AdamW(
            self.parameters(), lr=self.lr, weight_decay=1e-4)
        lr_scheduler = WarmupCosineSchedule(
            optimizer, warmup_steps=self.epoch/10, t_total=self.epoch)
        return [optimizer], [lr_scheduler]

    def training_step(self, batch, batch_idx):
        img, lab = batch
        pred = self.forward(img)
        loss_train = self.loss_fn(pred, lab)
        return {"loss": loss_train}

    def training_epoch_end(self, outputs):
        avg_loss = torch.stack([x["loss"] for x in outputs]).mean()
        self.log_dict({"loss": avg_loss, 'step': self.current_epoch})

    def validation_step(self, batch, batch_idx):
        img_val, lab_val = batch
        pred_val = self.forward(img_val)
        pred_val = self.post_trans(pred_val)
        self.dice(y_pred=pred_val, y=lab_val)

    def validation_epoch_end(self, outputs):
        metric = self.dice.aggregate().item()
        self.dice.reset()
        self.log_dict({"dice_val": metric, 'step': self.current_epoch})


MAX_EPOCH = 16

checkpoint_callback = ModelCheckpoint(
    dirpath='./checkPoint/', monitor='loss', save_top_k=2, mode='min', save_last=False)
early_stop_callback = EarlyStopping(
    monitor='dice_val', min_delta=0.00, patience=5, verbose=False, mode='max')


NET = SegResNet(dropout_prob=0.2, out_channels=1, blocks_down=[
                1, 2, 2, 4], blocks_up=[1, 1, 1], init_filters=16).to(cuda)
model = LitVnet(net=NET, loss_fn=DiceLoss(sigmoid=True),
                lr=1e-3, batch_size=32, img_size=64, epoch=MAX_EPOCH)

trainer = lit.Trainer(accelerator='auto', max_epochs=MAX_EPOCH, check_val_every_n_epoch=2, callbacks=[checkpoint_callback, early_stop_callback], log_every_n_steps=2,
                      amp_backend='apex', amp_level=2)

tuner = Tuner(trainer)
tuner.scale_batch_size(model, mode="binsearch")

---

LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
Batch size 2 succeeded, trying batch size 4
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
Batch size 4 succeeded, trying batch size 8
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
Batch size 8 succeeded, trying batch size 16
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
Batch size 16 succeeded, trying batch size 32
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
Batch size 32 succeeded, trying batch size 64
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
Batch size 64 failed, trying batch size 48
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
32

correct batch size should be 8

[akihironitta]

@ctwangtz Could you remind me of why the correct batch size should be 8?

[markvvw]

I tried

[rohitgr7]

it gives the most optimal batch size. Have you tried 32and why do you suggest it should be 8?