I’m training maskrcnn_resnet50_fpn and creating a Dataset as follows:
class CustomDataset(torch.utils.data.Dataset):
def __init__(self, root_dir):
self.root_dir = root_dir
self.image_files = os.listdir(os.path.join(root_dir, 'images'))
def __len__(self):
return len(self.image_files)
def __getitem__(self, idx):
img_name = os.path.join(self.root_dir, 'images', self.image_files[idx])
mask_name = os.path.join(self.root_dir, 'masks', self.image_files[idx])
#image = Image.open(img_name).convert('RGB').resize((1024,1024))
image = read_image(img_name, mode=ImageReadMode.RGB)
image = T.Resize((1024,1024))(image)
mask = Image.open(mask_name).convert('L').resize((1024,1024))
# Normalize pixel values to [0, 1]
image = torch.tensor(np.array(image)) / 255.0
mask = torch.tensor(np.array(mask)) / 255.0
mask = mask.unsqueeze(0)
# Apply horizontal flip with 50% probability
if torch.rand(1) < 0.5:
image = F.hflip(image)
mask = F.hflip(mask)
# Apply vertical flip with 50% probability
if torch.rand(1) < 0.5:
image = F.vflip(image)
mask = F.vflip(mask)
boxes = masks_to_boxes(mask)
labels = torch.tensor([1]) # Class label 1
target = {
'boxes': boxes,
'labels': labels,
'masks': mask
}
return image, target
The dataloaders are set up as follows:
train_dataset = CustomDataset(trainpath)
train_data_loader = DataLoader(train_dataset, batch_size=8, shuffle=True, collate_fn=collate_fn)
valid_dataset = CustomDataset(validatepath)
valid_data_loader = DataLoader(valid_dataset, batch_size=1, shuffle=True, collate_fn=collate_fn)
where collate_fn is defined as:
def collate_fn(batch):
return tuple(zip(*batch))
That works fine on the CPU. When I try to move the images and targets to the GPU using images, targets = images.to(device), targets.to(device), I get the following error:
---> 56 images, targets = images.to(device), targets.to(device)
AttributeError: 'tuple' object has no attribute 'to'
Trying to move the items separately,
images = images.to(device)
targets = targets.to(device)
throws the same error. The problem doesn’t seem to be that collate_fn is returning a tuple, and I’ve tried various tweaks to collate_fn but these throw different errors.
Advice appreciated.