[pytorch] ImageFolder 사용 방법

대용량 이미지 데이터를 데이터셋으로 관리하는데 편리한 모듈이 있다.

그 주인공은 Pytorch 에서 제공하는 ImageFolder다. 

 

 

이 글을 읽기 전 선수 지식 포스팅

• transforms.Compose 사용 방법

 

1) torchvision 설치

pip install torchvision

ImageFolder는 기본 torch가 아니라 torchvision을 통해 제공된다.

 

2) 모듈 import

from torchvision.datasets import ImageFolder

 

3) 사용 예제

 

예제 데이터구조

data 디렉토리를 기준으로 test, train  그 하위로  pizza, steak, sushi 디렉토리로 구성되어있다.

 

ImageFolder를 사용하기 위해 기본적인 디렉토리 골격을 지켜줘야 한다.

주목할 점은 제일 하위 디렉토리가 클래스명으로 이루어져 있다는 거다.

 

일반적으로 train, test 용으로 각각 ImageFolder객체를 생성한다.

train_imgfolder = ImageFolder(root='data/train',
                              transform=IMG_TRANSFORM, 
                              target_transform=None)   

test_imgfolder  = ImageFolder(root='data/test',
                              transform=IMG_TRANSFORM)

Argument	설명
root	클래스 디렉토리 상위 path 까지만 작성
transform	torchvision.transforms.Compose한 객체
target_transform	label은 transform False 지정

root에 입력된 path정보를 통해 해당경로 하위에 class명 디렉토리가 위치한 다는걸 ImageFolder가 알 수 있다.

 

이를 통해 클래스 정보를 다양한 형태로 얻을 수 있다.

print(train_imgfolder.classes)
# ['pizza', 'steak', 'sushi'])

print(train_imgfolder.class_to_idx)
# {'pizza': 0, 'steak': 1, 'sushi': 2}

 

DataLoader에 주입

train_dataloader = DataLoader(dataset=train_imgfolder, 
                              batch_size=32, 
                              num_workers=os.cpu_count(), 
                              shuffle=True)

test_dataloader  = DataLoader(dataset=test_imgfolder, 
                              batch_size=32, 
                              num_workers=os.cpu_count(), 
                              shuffle=False)

Argument	설명
dataset	ImageFolder 객체 인자로 넣어준다.
bach_size	배치 사이즈
num_workers	data loading할때 사용할 subprocess 숫자
shuffle	데이터 shuffle 여부(모델이 데이터 순서 학습 방지목적)

DataLoader에 train, test 데이터셋을 각각 넣어준다.

test의 경우 shuffle을 할 필요가 없으니 false로 지정하였다.

 

next(iter())로 batch별 shape를 확인할 수 있다.

batch_x, batch_y = next(iter(train_dataloader))

print(batch_x.shape, batch_y.shape) 
# torch.Size([32, 3, 64, 64]) torch.Size([32])