Dataloader for training

fastai
1

Hi,

I want to train a model on some .npz images and I Have created the following class for loading the data:

class data_gen(torch.utils.data.Dataset):
    def __init__(self, files):
        
        self.files = files
        
        
    def __getitem__(self, i):
        
        file1 = self.files[i]
        
        tmp = np.load(file1, allow_pickle=True)
        
        img = tmp['x']
        img = np.reshape(img,(1,img.shape[0], img.shape[1]))
        img = torch.from_numpy(img).float()
        
        return img

    def __len__(self): 
        
        return len(self.files)

And then, I create a DataLoader in the following way:

train_ds = data_gen(X_train)
test_ds = data_gen(X_test)
dls = DataLoaders.from_dsets(train_ds, test_ds, bs=batch_size, device='cuda:0')

However, when I get to training:

learn = Learner(dls,m,loss_func=F.mse_loss)
learn.fit(10,lr)

I get the following error:


RuntimeError: Expected 4-dimensional input for 4-dimensional weight [64, 1, 7, 7], but got 3-dimensional input of size [1, 512, 512] instead

I tried dls.train instead of dls (which is of the shape (16X1X512X512)) but even that doesn’t help.

I also tried to train on one batch of data but I would get the following error:

AttributeError: 'Tensor' object has no attribute 'train'

What should I do?

PS: My fastai version is 2.0.1