├── README.md └── main.py /README.md: -------------------------------------------------------------------------------- 1 | # Custom-CNN-Based-Classifier-in-PyTorch 2 | 3 | 본 튜토리얼에서는 **PyTorch** 를 사용하여 **Image Classifier** 를 만들어보도록 하겠습니다. 4 | 5 | 본 튜토리얼을 통해 다음 방법들을 배울 수 있습니다. 6 | 7 | * **CNN(Convolutional Neural Network)** 기반의 **Image Classifier** 모델 설계 방법 8 | * 기존의 데이터셋(MNIST, CIFAR-10 등)이 아닌 **Custom Dataset**(개인이 수집한 데이터셋)을 처리하기 위한 **PyTorch** 의 **Dataset** 및 **DataLoader** 사용 방법 9 | 10 | **본 튜토리얼에서는 PyTorch 의 Dataset 및 DataLoader 에 능숙해지기 위하여 PyTorch 의 ImageFolder 를 사용하지 않습니다.** 11 | 12 | ## Setup 13 | 14 | ### Requirements 15 | 16 | 본 튜토리얼에 앞서 다음의 라이브러리들을 필요로 합니다. 17 | * PyTorch 18 | * PIL 19 | 20 | 추가적으로 NVIDIA GPU 를 사용하신다면 CUDA 와 CUDNN 을 설치하는 것을 권장합니다. 21 | 22 | ### Custom Dataset 23 | 24 | 본 튜토리얼에서는 만화 **"쿠로코의 농구"** 의 등장 인물인 **"쿠로코"** 와 **"카가미"** 를 분류해보겠습니다. 25 | 26 | **Kuroko's basketball** 27 | 28 | ![60136-kuroko_no_basket-blue-basketball](https://user-images.githubusercontent.com/35001605/51658134-693e2700-1fea-11e9-8045-b2d49231246f.jpg) 29 | 30 | **Kuroko** 31 | 32 | ![65d739597a4f1bfd7085615cadf2c38e1367394101_full](https://user-images.githubusercontent.com/35001605/51658137-6c391780-1fea-11e9-9493-313b93f6166b.png) 33 | 34 | **Kagami** 35 | 36 | ![kurokonobaske_a03](https://user-images.githubusercontent.com/35001605/51658140-6e9b7180-1fea-11e9-9804-9ff222d3d079.jpg) 37 | 38 | #### Examples 39 | 40 | 정면 얼굴이 나온 사진 위주로 쿠로코 60장, 카가미 60 장 수집 41 | 42 | ![example_dataset3](https://user-images.githubusercontent.com/35001605/51650040-12285a00-1fca-11e9-95d1-189352ef2d58.PNG) 43 | 44 | ![example_dataset2](https://user-images.githubusercontent.com/35001605/51650039-105e9680-1fca-11e9-89be-868234ae3241.PNG) 45 | 46 | #### Download 47 | 48 | [다운로드 ](https://drive.google.com/open?id=1dQePxrd9xdtvLr9E-jiUb-TdyWG1EFlJ) 49 | 50 | 데이터셋은 학습을 위한 train 폴더와 성능 평가를 위한 test 폴더로 나누어져 있습니다. 51 | 52 | train 폴더와 test 폴더 내부에는 분류하고자 하는 클래스별로 하위 폴더가 존재하며 각 하위 폴더에는 해당 클래스에 속하는 이미지들이 저장되어 있습니다. 53 | 54 | ```python 55 | data/ 56 | train/ 57 | kuroko/ 58 | *.png 59 | kagami/ 60 | *.png 61 | test/ 62 | kuroko/ 63 | *.png 64 | kagami/ 65 | *.png 66 | ``` 67 | 68 | ### Structure of Directory 69 | 70 | 본 튜토리얼 에서는 **Custom Dataset**에 대한 처리를 보다 쉽게 하기 위하여 다음과 같이 main.py 파일(코드를 작성할 파일)과 **Custom Dataset**이 동일한 경로에 있음을 가정하겠습니다. 71 | 72 | ```python 73 | 74 | data/ 75 | train/ 76 | kuroko/ 77 | *.png 78 | kagami/ 79 | *.png 80 | test/ 81 | kuroko/ 82 | *.png 83 | kagami/ 84 | *.png 85 | 86 | main.py 87 | ``` 88 | ## Data Loading and Processing 89 | 90 | PyTorch 에는 데이터셋에 대한 처리를 용이하게 하기 위하여 **Dataset** 과 **DataLoader** 클래스를 제공합니다. 91 | 92 | **Dataset** 클래스는 torch.utils.data.Dataset 에 정의된 추상 클래스(Abstract class) 로써 사용자는 **Custom Dataset** 을 읽기 위하여 **Dataset** 클래스를 상속받는 클래스를 작성해야 합니다. 93 | 94 | **DataLoader**는 **Dataset** 클래스를 상속받는 클래스에 정의된 작업에 따라 데이터를 읽어오며, 이때 설정에 따라 원하는 **배치(Batch) 크기**로 데이터를 읽어올 수 있고 병렬 처리 설정, 데이터 셔플(Shuffle) 등의 작업을 설정할 수 있습니다. 95 | 96 | 본 튜토리얼에서는 **Dataset** 클래스를 상속받는 클래스를 **Dataset** 클래스라 칭하도록 하겠습니다. 97 | 98 | ### Dataset Class & DataLoader Class 99 | 100 | 일반적인 **Dataset** 클래스의 형태는 아래와 같습니다. 101 | 102 | ```python 103 | from torch.utils.data.dataset import Dataset 104 | 105 | class MyCustomDataset(Dataset): 106 | def __init__(self, ...): 107 | # stuff 108 | 109 | def __getitem__(self, index): 110 | # stuff 111 | return (img, label) 112 | 113 | def __len__(self): 114 | return count # of how many examples(images?) you have 115 | ``` 116 | **Custom Dataset** 을 읽기 위하여 다음의 3가지 함수를 정의해야 합니다. 117 | 118 | * `__init__()` 119 | 120 | * `__getitem__()` 121 | 122 | * `__len__()` 123 | 124 | `__init__()` 함수는 클래스 생성자로써 데이터에 대한 Transform(데이터 형 변환, Augmentation 등)을 설정하고 데이터를 읽기 위한 기초적인 초기화 작업들을 수행하도록 정의합니다. 125 | 126 | `__getitem__()` 함수는 **Custom Dataset** 에 존재하는 데이터를 읽고 반환하는 함수입니다. 따라서 본인이 어떤 작업을 수행하는지에 따라 반환하는 값들이 달라질 수 있습니다. 본 튜토리얼에서 구현할 작업은 **Image Classifier** 이므로 `__getitem__()` 함수는 이미지와 해당 이미지가 어떤 클래스에 속하는지에 대한 값을 반환할 것입니다. **1개의 데이터셋에 대한 값((이미지 , 클래스))만을 반환하면 됩니다!** 127 | 128 | 주의할 점은 `__getitem__()` 을 통해 반환되는 값이 PyTorch 에서 처리 가능한 데이터 타입(tensor, numpy array etc.)이 아닐 경우, **DataLoader** 를 통해 데이터를 읽을 때 다음과 같은 에러가 발생될 것입니다. 129 | 130 | `TypeError: batch must contain tensors, numbers, dicts or lists; found ` 131 | 132 | `__len__()` 함수는 데이터셋의 크기를 반환하는 함수입니다. **Image Classifier** 로 예를 들면, 우리가 가진 이미지의 갯수가 곧 데이터셋의 크기 입니다. 즉 50장을 가지고 있다면 `__len__()` 함수는 50 을 반환해야 합니다. 133 | 134 | ### Programming 135 | 136 | #### modules 137 | ```python 138 | import os 139 | from PIL import Image 140 | 141 | import torch 142 | from torch.utils.data import Dataset, DataLoader 143 | from torch import nn 144 | from torchvision import transforms 145 | ``` 146 | #### Declaration & Definition class 147 | ```python 148 | class CustomImageDataset(Dataset): 149 | def read_data_set(self): 150 | 151 | all_img_files = [] 152 | all_labels = [] 153 | 154 | class_names = os.walk(self.data_set_path).__next__()[1] 155 | 156 | for index, class_name in enumerate(class_names): 157 | label = index 158 | img_dir = os.path.join(self.data_set_path, class_name) 159 | img_files = os.walk(img_dir).__next__()[2] 160 | 161 | for img_file in img_files: 162 | img_file = os.path.join(img_dir, img_file) 163 | img = Image.open(img_file) 164 | if img is not None: 165 | all_img_files.append(img_file) 166 | all_labels.append(label) 167 | 168 | return all_img_files, all_labels, len(all_img_files), len(class_names) 169 | 170 | def __init__(self, data_set_path, transforms=None): 171 | self.data_set_path = data_set_path 172 | self.image_files_path, self.labels, self.length, self.num_classes = self.read_data_set() 173 | self.transforms = transforms 174 | 175 | def __getitem__(self, index): 176 | image = Image.open(self.image_files_path[index]) 177 | image = image.convert("RGB") 178 | 179 | if self.transforms is not None: 180 | image = self.transforms(image) 181 | 182 | return {'image': image, 'label': self.labels[index]} 183 | 184 | def __len__(self): 185 | return self.length 186 | ``` 187 | 188 | **Image Classifer**를 위한 **Dataset** 클래스는 위와 같이 구현되며 실제 객체 생성시에는 아래와 같이 사용됩니다. 189 | 190 | 191 | ```python 192 | custom_dataset_train = CustomImageDataset('./data/train',transforms=transfrom_train) 193 | custom_dataset_test = CustomImageDataset('./data/test',transforms=transfrom_test) 194 | ``` 195 | 196 | CustomImageDataset 클래스의 내부를 살펴보면 생성자(`__init__`)를 통하여 이미지들이 저장된 경로(data_set_path)를 입력받게 되고 이 값은 `self.data_set_path` 에 저장됩니다. 197 | 198 | ```python 199 | def __init__(self, data_set_path, transforms=None): 200 | ``` 201 | 202 | 이 후 `__init__` 함수에서는 `read_data_set` 함수를 호출하고 `self.data_set_path` 에 저장된 경로를 통해 이미지들의 경로(`self.image_files_path`), 각 이미지가 속한 클래스(`self.labels`), 데이터 셋의 크기(`self.length`), 분류하고자 하는 클래스의 갯수(`self.num_classes`) 를 초기화 합니다. 203 | -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | import os 2 | from PIL import Image 3 | 4 | import torch 5 | from torch.utils.data import Dataset, DataLoader 6 | from torch import nn 7 | from torchvision import transforms 8 | 9 | 10 | class CustomImageDataset(Dataset): 11 | def read_data_set(self): 12 | 13 | all_img_files = [] 14 | all_labels = [] 15 | 16 | class_names = os.walk(self.data_set_path).__next__()[1] 17 | 18 | for index, class_name in enumerate(class_names): 19 | label = index 20 | img_dir = os.path.join(self.data_set_path, class_name) 21 | img_files = os.walk(img_dir).__next__()[2] 22 | 23 | for img_file in img_files: 24 | img_file = os.path.join(img_dir, img_file) 25 | img = Image.open(img_file) 26 | if img is not None: 27 | all_img_files.append(img_file) 28 | all_labels.append(label) 29 | 30 | return all_img_files, all_labels, len(all_img_files), len(class_names) 31 | 32 | def __init__(self, data_set_path, transforms=None): 33 | self.data_set_path = data_set_path 34 | self.image_files_path, self.labels, self.length, self.num_classes = self.read_data_set() 35 | self.transforms = transforms 36 | 37 | def __getitem__(self, index): 38 | image = Image.open(self.image_files_path[index]) 39 | image = image.convert("RGB") 40 | 41 | if self.transforms is not None: 42 | image = self.transforms(image) 43 | 44 | return {'image': image, 'label': self.labels[index]} 45 | 46 | def __len__(self): 47 | return self.length 48 | 49 | 50 | class CustomConvNet(nn.Module): 51 | def __init__(self, num_classes): 52 | super(CustomConvNet, self).__init__() 53 | 54 | self.layer1 = self.conv_module(3, 16) 55 | self.layer2 = self.conv_module(16, 32) 56 | self.layer3 = self.conv_module(32, 64) 57 | self.layer4 = self.conv_module(64, 128) 58 | self.layer5 = self.conv_module(128, 256) 59 | self.gap = self.global_avg_pool(256, num_classes) 60 | 61 | def forward(self, x): 62 | out = self.layer1(x) 63 | out = self.layer2(out) 64 | out = self.layer3(out) 65 | out = self.layer4(out) 66 | out = self.layer5(out) 67 | out = self.gap(out) 68 | out = out.view(-1, num_classes) 69 | 70 | return out 71 | 72 | def conv_module(self, in_num, out_num): 73 | return nn.Sequential( 74 | nn.Conv2d(in_num, out_num, kernel_size=3, stride=1, padding=1), 75 | nn.BatchNorm2d(out_num), 76 | nn.LeakyReLU(), 77 | nn.MaxPool2d(kernel_size=2, stride=2)) 78 | 79 | def global_avg_pool(self, in_num, out_num): 80 | return nn.Sequential( 81 | nn.Conv2d(in_num, out_num, kernel_size=3, stride=1, padding=1), 82 | nn.BatchNorm2d(out_num), 83 | nn.LeakyReLU(), 84 | nn.AdaptiveAvgPool2d((1, 1))) 85 | 86 | 87 | hyper_param_epoch = 20 88 | hyper_param_batch = 8 89 | hyper_param_learning_rate = 0.001 90 | 91 | transforms_train = transforms.Compose([transforms.Resize((128, 128)), 92 | transforms.RandomRotation(10.), 93 | transforms.ToTensor()]) 94 | 95 | transforms_test = transforms.Compose([transforms.Resize((128, 128)), 96 | transforms.ToTensor()]) 97 | 98 | train_data_set = CustomImageDataset(data_set_path="./data/train", transforms=transforms_train) 99 | train_loader = DataLoader(train_data_set, batch_size=hyper_param_batch, shuffle=True) 100 | 101 | test_data_set = CustomImageDataset(data_set_path="./data/test", transforms=transforms_test) 102 | test_loader = DataLoader(test_data_set, batch_size=hyper_param_batch, shuffle=True) 103 | 104 | if not (train_data_set.num_classes == test_data_set.num_classes): 105 | print("error: Numbers of class in training set and test set are not equal") 106 | exit() 107 | 108 | device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') 109 | 110 | num_classes = train_data_set.num_classes 111 | custom_model = CustomConvNet(num_classes=num_classes).to(device) 112 | 113 | # Loss and optimizer 114 | criterion = nn.CrossEntropyLoss() 115 | optimizer = torch.optim.Adam(custom_model.parameters(), lr=hyper_param_learning_rate) 116 | 117 | for e in range(hyper_param_epoch): 118 | for i_batch, item in enumerate(train_loader): 119 | images = item['image'].to(device) 120 | labels = item['label'].to(device) 121 | 122 | # Forward pass 123 | outputs = custom_model(images) 124 | loss = criterion(outputs, labels) 125 | 126 | # Backward and optimize 127 | optimizer.zero_grad() 128 | loss.backward() 129 | optimizer.step() 130 | 131 | if (i_batch + 1) % hyper_param_batch == 0: 132 | print('Epoch [{}/{}], Loss: {:.4f}' 133 | .format(e + 1, hyper_param_epoch, loss.item())) 134 | 135 | # Test the model 136 | custom_model.eval() # eval mode (batchnorm uses moving mean/variance instead of mini-batch mean/variance) 137 | with torch.no_grad(): 138 | correct = 0 139 | total = 0 140 | for item in test_loader: 141 | images = item['image'].to(device) 142 | labels = item['label'].to(device) 143 | outputs = custom_model(images) 144 | _, predicted = torch.max(outputs.data, 1) 145 | total += len(labels) 146 | correct += (predicted == labels).sum().item() 147 | 148 | print('Test Accuracy of the model on the {} test images: {} %'.format(total, 100 * correct / total)) 149 | --------------------------------------------------------------------------------