S3 Torch-connector loads data as a list of tensors

[DavidMacha]

I'm setting up a model and about to start training a dataset from an S3 bucket. To load the data from S3 I'm using s3torchconnector.S3MapDataset.from_prefix which loads the data into the Sagemaker space I'm using. However, when I start to training, this S3torchconnector gives a list of tensors (or a Map dataset) instead of a tensor to be processed by samples = samples.to(device, non_blocking=True). As expected I'm getting the Error 'list' object has no attribute 'to'.

Here some snippets of my code:

def load_image(object):
   img = Image.open(object)
   resize = transforms.Resize(size=(224, 224))
   img = resize(img)
   img = transforms.functional.pil_to_tensor(img)
   return (object.key, torchvision.transforms.functional.convert_image_dtype(img, dtype=torch.float32))


dataset_train = s3torchconnector.S3MapDataset.from_prefix(
   args.IMAGES_URI, 
   region=args.REGION, 
   transform=load_image,
) # here apparently s3torchconnector is giving a list of tensors

.

data_loader_train = torch.utils.data.DataLoader(
    dataset_train, sampler=sampler_train,
    batch_size=args.batch_size,
    num_workers=args.num_workers,
    pin_memory=args.pin_mem,
    drop_last=True,
)

.

# define the model
model = models_mae.__dict__[args.model](norm_pix_loss=args.norm_pix_loss)

.

model.to(device)
for epoch in range(args.start_epoch, args.epochs):
    if args.distributed:
        data_loader_train.sampler.set_epoch(epoch)
    train_stats = train_one_epoch(
        model, data_loader_train,
        optimizer, device, epoch, loss_scaler,
        log_writer=log_writer,
        args=args
    )

This code works well with a custom dataset and local files in my Sagemaker space, but I cannot implement that way since I need to train millions of images, hence using an S3 dataloader. I'm not sure if S3torchconnector just works this way, or if it is possible to transform whatever it loads into a single Tensor to be processed by the Model and the sampler.

[Mattral]

The issue was resolved by adding a conditional structure to handle the type of the samples object before sending it to the device. Specifically, the following logic was used in the training loop (or wherever samples is processed):

if isinstance(samples, list):
    samples = samples[1].to(device, non_blocking=True)
else:
    samples = samples.to(device, non_blocking=True)

This approach ensures compatibility with both:

• Locally stored datasets in SageMaker, where samples is already a tensor.
• Datasets loaded using S3torchconnector.S3MapDataset, where samples is returned as a list (e.g., a tuple like (key, tensor)).

By checking the type and selecting the appropriate element (in this case, samples[1], which is the image tensor), the code is able to process the batch correctly regardless of the source of the dataset. This allows seamless training using either local or S3-based datasets.