How do I re-train Torchvision's keypoint R-CNN with a custom COCO-style dataset?

Question

I've created a custom COCO keypoints style dataset using COCO annotator and want to retrain Torchvision's Keypoint R-CNN on it. I'm trying to use torchvision's CocoDetection dataset class to load the data, and I had to rewrite the _load_image method because my dataset has subdirectories. I then tried to wrap the dataset in a dataloader and got the following error:

>>> dl = DataLoader(coco, batch_size=4)
>>> feat, lbl = next(iter(dl))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/sam/.local/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 521, in __next__
    data = self._next_data()
  File "/home/sam/.local/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 561, in _next_data
    data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
  File "/home/sam/.local/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py", line 47, in fetch
    return self.collate_fn(data)
  File "/home/sam/.local/lib/python3.8/site-packages/torch/utils/data/_utils/collate.py", line 84, in default_collate
    return [default_collate(samples) for samples in transposed]
  File "/home/sam/.local/lib/python3.8/site-packages/torch/utils/data/_utils/collate.py", line 84, in <listcomp>
    return [default_collate(samples) for samples in transposed]
  File "/home/sam/.local/lib/python3.8/site-packages/torch/utils/data/_utils/collate.py", line 84, in default_collate
    return [default_collate(samples) for samples in transposed]
  File "/home/sam/.local/lib/python3.8/site-packages/torch/utils/data/_utils/collate.py", line 84, in <listcomp>
    return [default_collate(samples) for samples in transposed]
  File "/home/sam/.local/lib/python3.8/site-packages/torch/utils/data/_utils/collate.py", line 74, in default_collate
    return {key: default_collate([d[key] for d in batch]) for key in elem}
  File "/home/sam/.local/lib/python3.8/site-packages/torch/utils/data/_utils/collate.py", line 74, in <dictcomp>
    return {key: default_collate([d[key] for d in batch]) for key in elem}
  File "/home/sam/.local/lib/python3.8/site-packages/torch/utils/data/_utils/collate.py", line 84, in default_collate
    return [default_collate(samples) for samples in transposed]
  File "/home/sam/.local/lib/python3.8/site-packages/torch/utils/data/_utils/collate.py", line 84, in <listcomp>
    return [default_collate(samples) for samples in transposed]
  File "/home/sam/.local/lib/python3.8/site-packages/torch/utils/data/_utils/collate.py", line 82, in default_collate
    raise RuntimeError('each element in list of batch should be of equal size')
RuntimeError: each element in list of batch should be of equal size

Is trying to put the dataset into a dataloader even the right thing to do, given that Keypoint R-CNN expects a list of [channels, height, width] tensors?

Additionally, I'm having trouble figuring out how I should actually train the model when I've got the data in an acceptable format. I've looked at https://pytorch.org/tutorials/beginner/transfer_learning_tutorial.html#convnet-as-fixed-feature-extractor and https://github.com/pytorch/vision/tree/master/references/detection and am still kind of confused. Can I get some directions on how to train the model on a machine with a single GPU?

How do I re-train Torchvision's keypoint R-CNN with a custom COCO-style dataset?

0 Answers0