How should monai.data.DataLoader() be passed to training iterator?

[relyativist]

#4545
From my issue I found that every get_item method of monai.data.DataLoader() call generates random transform. However, I
am bit confised how should monai.data.DataLoader() be passed in training loop as enumerate object to get same transformation for 'seg' and 'image'?

def one_epoch(model, 
                criterion, 
                opt, 
                config, 
                dataloader, 
                device, 
                writer, 
                epoch, 
                metric_dict_epoch, 
                n_iters_total=0,
                augmentation=None, 
                is_train=True):


    # use amp to accelerate training
    if config.opt.use_scaler:
        scaler = torch.cuda.amp.GradScaler()

    phase_name = 'train' if is_train else 'val'
    loss_name = config.opt.criterion
    metric_dict = defaultdict(list)
    target_metric_name = config.model.target_metric_name 

    if not is_train:
        model.eval()
    else:
        model.train()

    # used to turn on/off gradients
    grad_context = torch.autograd.enable_grad if is_train else torch.no_grad
    with grad_context():
        iterator = enumerate(dataloader)
        val_predictions = {}
        for i, data in iterator:
        for iter_i, data_tensors in iterator:
            brain_tensor = data_tensors['image']
            label_tensor = data_tensors['seg']
            mask_tensor = data_tensors['image'][0][2]

            print(brain_tensor.shape, label_tensor.shape)

            t1 = time.time()
        
            brain_tensor = brain_tensor.to(device)
            label_tensor = label_tensor.to(device)
            mask_tensor = mask_tensor.to(device)
            # forward pass
            with autocast(enabled=config.opt.use_scaler):
                label_tensor_predicted = model(brain_tensor) # -> [bs,1,ps,ps,ps]

                loss = criterion(label_tensor_predicted, label_tensor) 


            if is_train:
                opt.zero_grad()

                if config.opt.use_scaler:
                    scaler.scale(loss).backward()
                else:
                    loss.backward()
                
                if hasattr(config.opt, "grad_clip"):
                    if config.opt.use_scaler:
                        scaler.unscale_(opt)
                        torch.nn.utils.clip_grad_norm_(model.parameters(),
                                                           config.opt.grad_clip)

                metric_dict['grad_norm'].append(calc_gradient_norm(filter(lambda x: x[1].requires_grad, 
                                                model.named_parameters())))

                if config.opt.use_scaler:
                    scaler.step(opt)
                    scaler.update()
                else:
                    opt.step()


            t2 = time.time()    
            dt = t2-t1 # inference time
            
            metric_dict[f'batch_time'].append(dt)
            metric_dict[f'{loss_name}'].append(loss.item())
            label_tensor_predicted = label_tensor_predicted*mask_tensor
            dice_score = DiceScoreBinary(label_tensor_predicted, label_tensor)
            coverage = (label_tensor_predicted*label_tensor).sum() / label_tensor.sum()
            
            if not is_train:
#                #label = dataloader.dataset.labels[iter_i]
                 #label = dataloader.dataset[iter_i]['seg']
                label = dataloader.dataset.data[iter_i]['seg'].split('/')[6].split('.')[0]
                val_predictions[label] = label_tensor_predicted.detach().cpu().numpy()
            
            metric_dict['coverage'].append(coverage.item())
            metric_dict['dice_score'].append(dice_score.item())
            
            #########
            # PRINT #
            #########
            message = f'For {phase_name}, iter: {iter_i},'
            for title, value in metric_dict.items():
                if title == 'grad_norm':
                    v = np.round(value[-1],6)
                else:
                    v = np.round(value[-1],3)
                message+=f' {title}:{v}'
            print(message)

            # print(f'Epoch: {epoch}, Iter: {iter_i}, \n \
            # Loss_{loss_name}: {loss.item()}, Dice-score: {dice_score.item()}, \n \
            # time: {np.round(dt,2)}-s')

            if is_train and writer is not None:
                for title, value in metric_dict.items():
                    writer.add_scalar(f"{phase_name}_{title}", value[-1], n_iters_total)

            n_iters_total += 1

    target_metric = 0
    for title, value in metric_dict.items():
        m = np.mean(value)
        metric_dict_epoch[phase_name + '_' + title].append(m)
        if title == target_metric_name:
            target_metric = m
        if writer is not None:
            writer.add_scalar(f"{phase_name}_{title}_epoch", m, epoch)
            
    #####################
    # SAVING BEST PREDS #
    #####################
    target_metrics_epoch = metric_dict_epoch[f'val_{target_metric_name}']
    if not is_train:
        if config.dataset.save_best_val_predictions:
            # use greedy-saving: save only if the target metric is improved
            if len(target_metrics_epoch) == 1 or target_metrics_epoch[-1] >= target_metrics_epoch[-2]:
                for label, pred in val_predictions.items():
                        #torch.save(pred, os.path.join(config.dataset.val_preds_path, f'{label}'))
                        torch.save(pred, os.path.join(config.dataset.val_preds_path, f'{label}'))


    return n_iters_total, target_metric

I create input dataloader with folowwing steps:
data = {'image': ['/tmp/sub-1/sub-1_t1.nii.gz', '/tmp/sub-t2.nii.gz'], 'seg': '/tmp/1.nii.gz'}

train_transf = Compose(
    [
        LoadImaged(keys=["image", "seg"]),
        EnsureChannelFirstd(keys=["image", "seg"]),
        RandRotated(keys=["image", "seg"], range_x=0.0, range_y=0.0, range_z=0.75, prob=1),
        Resized(keys=["image", "seg"], spatial_size=(128,128,128), mode=('area')),
        EnsureTyped(keys=["image", "seg"], dtype=torch.float),
    ]
)

train_ds = monai.data.Dataset(data=train_files, transform=train_transf)
train_loader = DataLoader(
    train_ds,
    batch_size=2,
    shuffle=False,
    num_workers=0,
    collate_fn=list_data_collate,
    pin_memory=torch.cuda.is_available(),
)

When I pass train_loader in one_epoch and print it with:
print(brain_tensor.shape, label_tensor.shape) I got the dimension equal to torch.Size([2, 3, 128, 128, 128]) torch.Size([2, 1, 128, 128, 128])

Does it right way to load it to iterator?