[Pytorch] ImageFolder label기준으로 split하기

ResNet을 구현한 다음 학습을 시켜보았는데 한 에포크를 도는데 5분쯤 걸렸다.

단순 테스트용으로 돌리는건데 너무 시간이 길다고 생각해서 라벨을 쪼개려고 했는데

 

처음에는 Stratified한 방식으로 쪼개려 했는데 생각해보니 데이터셋으로 사용한  imagenet-mini는 이미 class별 이미지가 너무 적다 생각하여 label을 index로 지정해서 100개 정도만 뽑아서 데이터셋으로 만드려했다.

 

로컬폴더를 900개 제거하면 ImageFolder가 알아서 해주겠지만 뭔가 데이터셋을 제거한다는 방식이 세련되 보이지도 않았고 이후에도 train과 test를 쪼갤일이 많을 텐데 미리 연습 좀 해보자는 느낌으로 방법들을 서칭했다.

별거 아닌 고집이었는데 구현하는데 시간이 예상보다 오래 걸렸다.

 

이번에는 내가 ImageFolder를 사용해서 label을 index로 사용해 Dataset을 split했던 방법을 써보겠다.

 

 

Subset

Subset을 이용하여 Dataset을 split하기로 결정했다. 고로 먼저 Subset에 대해 간단히 설명하겠다.

Dataset과 그로부터 뽑아내고 싶은 index들을 넣어주면 그 index만 가지는 Dataset을 반환해준다.

정확히는 Dataset이 아니라 Dataset으로부터 파생된 Subset을 반환하는데 Dataloader로 넣어주면 원하는대로 동작한다.

 

 

train = datasets.ImageFolder(root='../../data/imagenet-mini/train/', transform=transform)
temp_subset = torch.utils.data.Subset(train, indices=[0,1,2])

print(train)
print(temp_subset)
print(temp_subset.dataset)

Dataset ImageFolder
    Number of datapoints: 34745
    Root location: ../../data/imagenet-mini/train/
    StandardTransform
Transform: Compose(
               ToTensor()
               Resize(size=[224, 224], interpolation=bilinear, max_size=None, antialias=True)
           )
           
           
<torch.utils.data.dataset.Subset at 0x7f4a8bdaf8e0>


Dataset ImageFolder
    Number of datapoints: 34745
    Root location: ../../data/imagenet-mini/train/
    StandardTransform
Transform: Compose(
               ToTensor()
               Resize(size=[224, 224], interpolation=bilinear, max_size=None, antialias=True)
           )

Subset은 기존에 생성했던 Dataset으로 파생되었고 이에 대한 정보를 가지고 있다.

 

 

print(len(train))
print(len(temp_subset))

34745
3

넣어준 인덱스가 [0, 1, 2]이므로 temp_subset은 3의 길이를 가지고 있음을 확인할 수 있다.

 

 

 

Split Dataset using Subset

random_indices = np.random.choice(range(1000),500, replace=False)
temp=[]
for inx, label in enumerate(train.targets):
    if label in random_indices:
        temp.append(inx)
    else:
        pass
temp_test=[]
for inx, label in enumerate(test.targets):
    if label in random_indices:
        temp_test.append(inx)
    else:
        pass
        
        
subset = torch.utils.data.Subset(train, indices=temp)
subset_test = torch.utils.data.Subset(test, indices=temp_test)

위의 코드는 기존의 1000개의 라벨을 가지고 있던 train과 test에서 500개의 random한 label값을 추출한 다음 train과 test Dataset의 라벨값을 enumerating하면 조건에 맞는 index를 뽑는다.

그 인덱스를 Subset에 넣어주면 index별로 쪼개진 Dataset이 만들어진다.

 

 

print(len(train))
print(len(subset))

34745
17459

데이터셋 길이를 통해 약 절반쯤으로 데이터셋이 쪼개진걸 확인할 수 있다.

 

 

주의점: 이후 train과 test를 쪼개며 확인했는데 train과 test를 같은 dataset에서 각기 다른 index로 Subset을 통해 생성하면 둘이 메모리를 공유하게 되며 train에 적용하는 증강이 test에도 적용되는 등의 문제가 생긴다.

따라서 이런 경우에는 copy를 사용하여 각기 다른 메모리의 데이터셋에 Subset을 적용해야한다.

---

Tensorflow에서는 쉽게 하던 일이었는데 Pytorch로 넘어오며 예상 외의 부분에서 난관에 부딪히는 느낌이다.

서칭하고 구현하는데 1시간은 훌쩍 넘게 걸린 느낌인데 아직까지도 Pytorch가 손에 익지 않았다고 다시 느꼈다.

부캠에서 경진대회를하며 이러한 문제가 발생했는데 이런 기초적인 문제가 아직까지 파이토치에 남아있을것 같지는 않고  아마 버전업을 하거나 Subset클래스를 만들면서 생긴 일시적 에러가 아닐까 생각한다.

 

출처: https://pytorch.org/docs/stable/data.html?highlight=subset#torch.utils.data.Subset