[Pytorch] ImageNet Dataset 사용하기

안녕하세요. 오늘은 영상 분류 관련 논문에서 필수적으로 사용되는 데이터셋인 ImageNet에 대해서 알아보고 pytorch를 이용해서 DataLoader를 만들어보도록 하겠습니다.

 

1. ImageNet Dataset 소개

ImageNet Dataset은 MNIST, CIFAR, STL, SVHN 등과 더불어 같이 영상 분류에 있어 가장 유명한 데이터셋 중 하나 입니다. 하지만, ImageNet이 다른 영상 분류 데이터셋보다 더욱 유명한 이유는 데이터셋의 규모가 비교도 안될 정도로 거대하기 때문입니다. 

 

- MNIST : 0 ~ 9의 손글씨 $28 \times 28$ 이미지 데이터로 학습 60,000개, 시험 10,000개로 구성

- CIFAR-10/100 : 일반적인 객체 10개 또는 100개의 클래스로 이루어진 $32 \times 32$ 이미지 데이터로 학습 50,000개, 시험 10,000개로 구성

- STL-10 : 일반적인 객체 10개의 클래스로 이루어진 $96 \times 96$ 이미지 데이터로 학습 5,000개, 시험 8,000개로 구성

- SVHN : 구글의 Street View를 이용해서 가정집의 우편번호 $32 \times 32$ 이미지 데이터로 학습 73,257개, 시험 26,032개로 구성

 

위 데이터셋을 보시면 아무리 많아도 100,000개가 되지 않습니다. 또한, 다양한 실험을 할 때 빨리 진행되어야 가설 검증 및 재실험을 사용할 수 있기 때문에 주로 CIFAR-10/100 또는 STL-10 데이터셋을 사용하는 편 입니다. 하지만, 이러한 데이터셋들의 문제점은 현실세계에 바로 적용하기에는 어려운 문제가 있습니다. 바로 제한적인 클래스의 개수와 해상도, 그리고 데이터셋의 개수 입니다. 이 세상에는 수많은 종류의 객체가 존재합니다. 그리고 그 객체 안에서도 다양한 종류로 나뉘게 되죠. 당장, 물고기만 예를 들어도 구피, 금붕어, 잉어, ... 굉장히 많은 종류가 있죠. 이러한 이유로 ImageNet은 필수적인 데이터셋이죠. 

 

ImageNet은 다양한 해상도를 가지는 영상이 학습 1,200,000개, 검증 50,000개가 넘는 영상이 각각 1,000개의 클래스로 구성된 데이터셋입니다. 기본적으로 다른 데이터셋에 비해서 그 규모가 엄청난 것을 볼 수 있죠. 그런데 여기서 시험 데이터셋에 대한 이야기를 하지 않는 이유는 여전히 시험 데이터셋이 공개되지 않고 있으며 자체 평가 홈페이지에 업로드해야 시험 데이터셋으로 평가된 결과를 얻을 수 있습니다. 따라서, 일반적으로 논문을 작성할 때는 학습 데이터셋 (학습용)과 검증 데이터셋 (시험용)만을 주로 사용합니다. 

 

2. ImageNet Dataset 다운로드

 

ImageNet

 

ImageNet Dataset은 위 링크에서 로그인 후 직접 다운로드 가능합니다. 주의할 점은 gmail이나 hotmail과 같은 도메인은 사용할 수 없으며 학교 및 기관 이메일을 사용하셔야 회원가입을 진행할 수 있습니다. 

 

 

회원가능 후 로그인을 하시면 Download 버튼을 누를 수 있게 됩니다. 그러면 위와 같은 화면이 뜰텐데요. 여기서 ILSVRC 메뉴에서 원하는 연도의 ImageNet 데이터셋을 다운로드 받을 수 있습니다. 저는 ImageNet2012를 다운받도록 하죠. 

 

 

그러면 위 화면을 볼 수 있습니다. 여기서, Task1, 2, 3로 나누어져 데이터셋이 제공되는 데, 저희는 영상 분류를 진행할 예정이기 때문에 Task1 데이터셋을 다운받으면 됩니다. 따라서, Images 메뉴의 Training images (Task 1 & 2)와 Validation images (all tasks)를 선택하고 다운로드 받으시면 됩니다. 그러면 다운로드가 진행되는데 각 데이터셋은 훈련 데이터는 138GB, 검증 데이터는 728MB이기 때문에 전부 다운로드 받는데 1~3일 정도 걸리게 됩니다. 저는 2일 정도 걸렸으니 네트워크 상태 확인 꼭 하시고 다운로드 하시길 바랍니다. (저는 중간에 끊겨서 몇 번 다시 다운받았습니다...)

 

다운로드하시면 2가지 tar 파일을 얻을 수 있습니다. 이제 각 압축파일을 원하는 위치에 옮기신 후 다음 명령어를 치시면 됩니다. 

 

[AI dataset] ImageNet Dataset 다운로드

 

# extract train data
mkdir train && mv ILSVRC2012_img_train.tar train/ && cd train
tar -xvf ILSVRC2012_img_train.tar && rm -f ILSVRC2012_img_train.tar
find . -name "*.tar" | while read NAME ; do mkdir -p "${NAME%.tar}"; tar -xvf "${NAME}" -C "${NAME%.tar}"; rm -f "${NAME}"; done
# extract validation data
cd ../ && mkdir val && mv ILSVRC2012_img_val.tar val/ && cd val && tar -xvf ILSVRC2012_img_val.tar
wget -qO- https://raw.githubusercontent.com/soumith/imagenetloader.torch/master/valprep.sh | bash

 

 

그러면 위 사진과 같은 폴더들이 압축해제가 됩니다. 여기서 폴더명은 클래스의 이름을 의미하는 데 아마 "n"으로 시작하게 될 겁니다. 이를 알아보기 쉽게 보기 위해 구글에 imagenet_class_index.json 파일까지 다운받아주시면 학습할 준비가 전부 끝나게 됩니다. 

 

3. ImageNet DataLoader

 

ImageNet DataLoader를 구현한 전체 코드는 제 깃허브에 있습니다. 

 

GitHub - skawngus1111/Study

 

기본적으로 각 영상은 1,000개의 폴더에 분류되어 저장되어있습니다. 1,000개의 클래스와 1,000개의 폴더... 맞습니다. 각 폴더 = 클래스를 의미하죠. 따라서, 폴더에 따른 클래스를 바로 넘겨주면 됩니다. 이를 위한 pytorch 함수가 존재하죠. 바로, ImageFolder입니다. 

 

train_dataset = datasets.ImageFolder(train_dir, transform)
val_dataset = datasets.ImageFolder(val_dir, transform)

 

훈련 및 검증 데이터가 존재하는 경로와 transform (Resize, Flipping, ... 등)을 넘겨주면 자동으로 저희가 필요했던 train_dataset/val_dataset이 만들어지게 됩니다. dataset이 정의되었으며 바로 DataLoader 함수로 넘겨주면 바로 사용할 수 있죠. 

 

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=2, shuffle=True)
val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=2, shuffle=False)

 

그런데, 아까도 말씀드렸다싶이 바로 클래스 이름을 해석하기에는 어렵기 때문에 imagenet_class_index.json 파일이 필요합니다. 이를 이용하면 쉽게 어떤 클래스인지 알 수 있죠. 

 

idx2label = []
cls2label = {}
with open(json_file, "r") as read_file:
    class_idx = json.load(read_file)
    idx2label = [class_idx[str(k)][1] for k in range(len(class_idx))]
    cls2label = {class_idx[str(k)][0]: class_idx[str(k)][1] for k in range(len(class_idx))}

 

여기까지 오셨으면 ImageNet Dataset을 사용하기 위한 모든 준비가 끝났습니다. 남은 것은 모델, 최적화 함수, 손실 함수를 정의하고 학습하는 것이죠. 이는 나중 포스팅에 모델을 학습할 때 설명드리도록 하겠습니다. 마무리로 ImageNet Dataset의 시각화를 진행하도록 하겠습니다. 

 

Granny_Smith

 

Brittany_spaniel

 

vine_snake