[Problem] problem occured when trained on custom dataset

[CaptainEven]

Hi, Thanks for your excellent work!
I want to train on my own dataset which consist of many different sub-dir paths, so warite a PyTorch Dataset with input of train.txt(a list of img paths from different sub-dir paths) as beblow:

class DatasetFromTxtList(Dataset):
    def __init__(self, txt_path):
        """
        Read data path from a TXT list file
        """
        if not os.path.isfile(txt_path):
            print("[Err]: invalid txt file path.")
            exit(-1)

        self.img_paths = []
        with open(txt_path, "r", encoding="utf-8") as f:
            for line in f.readlines():
                img_path = line.strip()
                self.img_paths.append(img_path)
        print("Total {:d} images found.".format(len(self.img_paths)))

        ## Define transformations
        self.normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                              std=[0.229, 0.224, 0.225])

        # MoCo v2's aug: similar to SimCLR https://arxiv.org/abs/2002.05709
        augmentations = [
            transforms.RandomResizedCrop(224, scale=(0.2, 1.0)),
            transforms.RandomApply([
                transforms.ColorJitter(0.4, 0.4, 0.4, 0.1)  # not strengthened
            ], p=0.8),
            transforms.RandomGrayscale(p=0.2),
            transforms.RandomApply([GaussianBlur([0.1, 2.0])], p=0.5),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            self.normalize
        ]

        self.T = transforms.Compose(augmentations)

    def __getitem__(self, idx):
        """
        """
        img_path = self.img_paths[idx]
        x = Image.open(img_path)

        q = self.T(x)
        k = self.T(x)

        return [q, k]

    def __len__(self):
        """
        """
        return len(self.img_paths)

and, i replace the train_dataset definition with:

    ## ----- Using customized dataset: reading sample from a txt list file...
    train_dataset = DatasetFromTxtList(args.train_txt)

instead of:

    train_dataset = datasets.ImageFolder(
        train_dir,
        simsiam.loader.TwoCropsTransform(transforms.Compose(augmentation)))

error is as follows:

Total 502335 images found.
Traceback (most recent call last):
  File "/mnt/diskb/even/SimSiam/my_simsiam.py", line 468, in <module>
    main()
  File "/mnt/diskb/even/SimSiam/my_simsiam.py", line 203, in main
    mp.spawn(main_worker, nprocs=n_gpus_per_node, args=(n_gpus_per_node, args))
  File "/usr/local/lib/python3.7/dist-packages/torch/multiprocessing/spawn.py", line 230, in spawn
    return start_processes(fn, args, nprocs, join, daemon, start_method='spawn')
  File "/usr/local/lib/python3.7/dist-packages/torch/multiprocessing/spawn.py", line 188, in start_processes
    while not context.join():
  File "/usr/local/lib/python3.7/dist-packages/torch/multiprocessing/spawn.py", line 150, in join
    raise ProcessRaisedException(msg, error_index, failed_process.pid)
torch.multiprocessing.spawn.ProcessRaisedException: 

-- Process 3 terminated with the following error:
Traceback (most recent call last):
  File "/usr/local/lib/python3.7/dist-packages/torch/multiprocessing/spawn.py", line 59, in _wrap
    fn(i, *args)
  File "/mnt/diskb/even/SimSiam/my_simsiam.py", line 354, in main_worker
    train(train_loader, model, criterion, optimizer, epoch, args)
  File "/mnt/diskb/even/SimSiam/my_simsiam.py", line 391, in train
    p1, p2, z1, z2 = model(x1=images[0], x2=images[1])
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/parallel/distributed.py", line 799, in forward
    output = self.module(*inputs[0], **kwargs[0])
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/mnt/diskb/even/SimSiam/simsiam/builder.py", line 55, in forward
    z1 = self.encoder(x1) # NxC
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/torchvision/models/resnet.py", line 249, in forward
    return self._forward_impl(x)
  File "/usr/local/lib/python3.7/dist-packages/torchvision/models/resnet.py", line 232, in _forward_impl
    x = self.conv1(x)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/conv.py", line 443, in forward
    return self._conv_forward(input, self.weight, self.bias)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/conv.py", line 440, in _conv_forward
    self.padding, self.dilation, self.groups)
RuntimeError: Expected 4-dimensional input for 4-dimensional weight [64, 3, 7, 7], but got 3-dimensional input of size [3, 224, 224] instead

How to solve this?

[CaptainEven]

The problem solved by return a dict from custom Dataset.

[CloudWalking0]

The problem solved by return a dict from custom Dataset.

Hello, I have met the same problem when trained on custom dataset. Would you mind providing more info about this issue?

Thank you very much~

[CloudWalking0]

The problem solved by return a dict from custom Dataset.

Hello, I have met the same problem when trained on custom dataset. Would you mind providing more info about this issue?

Thank you very much~
The problem solved by replace "for i, (images, _) in enumerate(train_loader):" with "for i, images in enumerate(train_loader):"

[wmrenr]

what should I do if trained on custom dataset ?