HW3_train.py

# -*- coding: utf-8 -*-
"""CNN_predict_ssh.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1c1N7jjIJuamySLG_iobRHdRLmwn-s1Ww
"""



from torchsummary import summary
import os
import numpy as np
import cv2
import torch
import torch.nn as nn
import torchvision.transforms as transforms
import pandas as pd
from torch.utils.data import DataLoader, Dataset
import time
import sys

"""#Read image
利用 OpenCV (cv2) 讀入照片並存放在 numpy array 中
"""

def readfile(path, label):
    # label 是一個 boolean variable，代表需不需要回傳 y 值
    image_dir = sorted(os.listdir(path))
    x = np.zeros((len(image_dir), 128, 128, 3), dtype=np.uint8)
    y = np.zeros((len(image_dir)), dtype=np.uint8)
    for i, file in enumerate(image_dir):
        img = cv2.imread(os.path.join(path, file))
        x[i, :, :] = cv2.resize(img,(128, 128))
        if label:
          y[i] = int(file.split("_")[0])
    if label:
      return x, y
    else:
      return x

print("Reading data")
train_x = np.load("train_x.npy")
train_y = np.load("train_y.npy")
print("Size of training data = {}".format(len(train_x)))
val_x = np.load("val_x.npy")
val_y = np.load("val_y.npy")

#分別將 training set、validation set、testing set 用 readfile 函式讀進來
'''
workspace_dir = sys.argv[1]
MODLE_PATH = "model_params/best_model3.pt"
print("Reading data")
test_x = readfile(os.path.join(workspace_dir, "testing"), False)
print("Size of Testing data = {}".format(len(test_x)))
'''
workspace_dir = sys.argv[1]
MODLE_PATH = 'model_best.pt'
test_x = np.load('test_x.npy')

print("Loading mean")
my_mean = np.load("mean.npy")
print("mean loaded")
print("Loading std")
my_std = np.load("std.npy")
print("std loaded")

#training 時做 data augmentation
train_transform = transforms.Compose([
    transforms.ToPILImage(),
    transforms.RandomHorizontalFlip(), #隨機將圖片水平翻轉
    transforms.RandomRotation(15), #隨機旋轉圖片
    transforms.ToTensor(), #將圖片轉成 Tensor，並把數值normalize到[0,1](data normalization)
    transforms.Normalize(my_mean,my_std),
])
#testing 時不需做 data augmentation
test_transform = transforms.Compose([
    transforms.ToPILImage(),                                    
    transforms.ToTensor(),
    transforms.Normalize(my_mean,my_std),
])
class ImgDataset(Dataset):
    def __init__(self, x, y=None, transform=None):
        self.x = x
        # label is required to be a LongTensor
        self.y = y
        if y is not None:
            self.y = torch.LongTensor(y)
        self.transform = transform
    def __len__(self):
        return len(self.x)
    def __getitem__(self, index):
        X = self.x[index]
        if self.transform is not None:
            X = self.transform(X)
        if self.y is not None:
            Y = self.y[index]
            return X, Y
        else:
            return X

batch_size = 128#128

"""# Model"""




train_val_x = np.concatenate((train_x, val_x), axis=0)
train_val_y = np.concatenate((train_y, val_y), axis=0)
train_val_set = ImgDataset(train_val_x, train_val_y, train_transform)
#train_val_loader = DataLoader(train_val_set, batch_size=batch_size, shuffle=True)
class Classifier(nn.Module):
    def __init__(self):
        super(Classifier, self).__init__()
        #torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding)
        #torch.nn.MaxPool2d(kernel_size, stride, padding)
        #input 維度 [3, 128, 128]
        self.cnn = nn.Sequential(
            nn.Conv2d(3, 64, 3, 1, 1),  # [64, 128, 128]
            nn.BatchNorm2d(64),
            nn.LeakyReLU(0.04),
            nn.MaxPool2d(2, 2, 0),      # [64, 64, 64]
            nn.Dropout(0.2),

            nn.Conv2d(64, 128, 3, 1, 1), # [128, 64, 64]
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.04),
            nn.MaxPool2d(2, 2, 0),      # [128, 32, 32]
            nn.Dropout(0.2),

            nn.Conv2d(128, 256, 3, 1, 1), # [256, 32, 32]
            nn.BatchNorm2d(256),
            nn.LeakyReLU(0.04),
            nn.MaxPool2d(2, 2, 0),      # [256, 16, 16]
            nn.Dropout(0.2),

            nn.Conv2d(256, 512, 3, 1, 1), # [512, 16, 16]
            nn.BatchNorm2d(512),
            nn.LeakyReLU(0.04),
            nn.MaxPool2d(2, 2, 0),       # [512, 8, 8]
            nn.Dropout(0.2),
            
            nn.Conv2d(512, 512, 3, 1, 1), # [512, 8, 8]
            nn.BatchNorm2d(512),
            nn.LeakyReLU(0.04),
            nn.MaxPool2d(2, 2, 0),       # [512, 4, 4]
            nn.Dropout(0.2),

            nn.Conv2d(512, 512, 3, 1, 1), # [512, 4, 4]
            nn.BatchNorm2d(512),
            nn.LeakyReLU(0.04),
            nn.MaxPool2d(2, 2, 0),       # [512, 2, 2]
            nn.Dropout(0.2),

            nn.Conv2d(512, 512, 3, 1, 1), # [512, 4, 4]
            nn.BatchNorm2d(512),
            nn.LeakyReLU(0.04),
            nn.MaxPool2d(2, 2, 0),       # [512, 1, 1]
            nn.Dropout(0.2),

            
        )
        self.fc = nn.Sequential(
            # nn.Linear(512, 1024),
            # #nn.Dropout(0.1),
            # nn.LeakyReLU(0.03),
            # nn.Linear(1024, 512),
            # #nn.Dropout(0.1),
            # nn.LeakyReLU(0.03),
            # nn.Linear(512, 256),
            # nn.Dropout(0.2),
            # nn.LeakyReLU(0.03),
            # nn.Linear(256, 128),
            # nn.Dropout(0.2),
            # nn.LeakyReLU(0.03),
            # nn.Linear(128, 11)
            nn.Linear(512,256), #nn.linear(dim of input, dim of output)
            nn.ReLU(),
            nn.Linear(256,128),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(128, 11)
        )

    def forward(self, x):
        out = self.cnn(x)
        out = out.view(out.size()[0], -1)
        return self.fc(out)
train_val_loader = torch.load('train_val_loader.pth')
# Train
model_best = Classifier().cuda()
summary(model_best,(3,128,128))
loss = nn.CrossEntropyLoss() # 因為是 classification task，所以 loss 使用 CrossEntropyLoss
optimizer = torch.optim.Adam(model_best.parameters(), lr=0.001,weight_decay=0.0005) # optimizer 使用 Adam
num_epoch = 100

for epoch in range(num_epoch):
    epoch_start_time = time.time()
    train_acc = 0.0
    train_loss = 0.0

    model_best.train()
    for i, data in enumerate(train_val_loader):
        optimizer.zero_grad()
        train_pred = model_best(data[0].cuda())
        batch_loss = loss(train_pred, data[1].cuda())
        batch_loss.backward()
        optimizer.step()

        train_acc += np.sum(np.argmax(train_pred.cpu().data.numpy(), axis=1) == data[1].numpy())
        train_loss += batch_loss.item()

        #將結果 print 出來
    print('[%03d/%03d] %2.2f sec(s) Train Acc: %3.6f Loss: %3.6f' % \
      (epoch + 1, num_epoch, time.time()-epoch_start_time, \
      train_acc/train_val_set.__len__(), train_loss/train_val_set.__len__()))



print("Saving model...")
torch.save(model_best.state_dict(), "model_best.pt")
print("model_best.pt saved")