Resolving Dataset Format, Variable Naming, and Data Loading Issues

[codebymov]

Hi,
I encountered and resolved several issues while working with the EfficientAD implementation. I'd like to share these fixes to help improve the codebase.

Issues Found:

1. ImageNetDataset return format mismatch causing tensor dimension errors
2. Inconsistent variable naming (train_iter vs iteration) in distillation training
3. Data loading format inconsistency between teacher and student training

Changes Made to Resolve:

1. Modified ImageNetDataset to return consistent dictionary format:

def __getitem__(self, idx):
    img, _ = self.dataset[idx]
    return {'image': img}

Fixed variable naming in distillation_training.py:

# Changed self.train_iter to self.iteration for consistency
scheduler = torch.optim.lr_scheduler.StepLR(
    optimizer, step_size=int(0.95 * self.iteration), gamma=0.1)

Updated data loading to handle dictionary format:

batch_sample = next(dataloader)['image'].cuda()

In the last step, replace the following:

s_imagenet_out = student(image_p[0].cuda())

with:

s_imagenet_out = student(image_p['image'].cuda())

These changes allow the training pipeline to run successfully while maintaining code consistency. I've tested these modifications with both the distillation training and student training phases