[PyTorch] Dataset을 Concat 할 수 있다고...?

Crowd Counting 목적으로 구축된  ShanghaiTech 데이터셋을 보면 A part/B part 로 데이터셋이 나누어져있다.

 

이는 목적은 같지만 데이터의 특징이 달라서 그런건데,  A, B를 합쳐서 학습을 해보고 싶다는 생각이 들었다.

 

기존에 데이터셋 클래스를, A part 따로 B part 따로 읽을 수 있게 구현해놨는데 이걸 합쳐서 읽을 수 있게 내가 직접 구현할까 하다가 한 번 찾아봤다.

 

역시나 있더라...

 

https://stackoverflow.com/a/67163131

Pytorch - Concatenating Datasets before using Dataloader

 

torch.utils.data.ConcatDataset 클래스를 이용하면 된다.

예시

    train_dataset_a = ds.ShanghaiCrowdDataset(root="ShanghaiTech", part="A", split="train", apply_augmentation=True)
    train_dataset_b = ds.ShanghaiCrowdDataset(root="ShanghaiTech", part="B", split="train", apply_augmentation=True)
    train_dataset = torch.utils.data.ConcatDataset([train_dataset_a, train_dataset_b])

 

더 놀라운건 아래처럼 써도 된다.

 

    train_dataset_a = ds.ShanghaiCrowdDataset(root="ShanghaiTech", part="A", split="train", apply_augmentation=True)
    train_dataset_b = ds.ShanghaiCrowdDataset(root="ShanghaiTech", part="B", split="train", apply_augmentation=True)
    train_dataset = train_dataset_a + train_dataset_b