import gc
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
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) # 创建指定维度全0的ndarray用来存放图片，彩色，3通道，把所有图片转化为128 * 128，
y = np.zeros((len(image_dir)), dtype=np.uint8) # 存放对应图片的标签，只有 0 ~ 10 所以用8位无符号数够了
for i, file in enumerate(image_dir):
img = cv2.imread(os.path.join(path, file)) # 拼接路径，目录 + 名字 读入一张图片
x[i, :, :] = cv2.resize(img, (128, 128)) # 对每张图片调整为128 * 128 同时保存在x的指定位置中
if label:
y[i] = int(file.split("_")[0]) # 0_0.jpg 以_来切分，取出第一个元素，即图片的标签
if label:
return x, y
else:
return x
# training 时候做 data augmentation（数据增强）
train_transform = transforms.Compose([
transforms.ToPILImage(), # 因为之前的类型是ndarray，不能进行后续的转化呢
transforms.RandomHorizontalFlip(), # 几率为0.5镜像
transforms.RandomRotation(15), # 0~15度的范围内旋转图片
transforms.ToTensor(), # 將图片成 Tensor，并 normalize 到 [0,1] (data normalization)，因为后面处理时候只吃tensor类型的数据
])
# testing 不做 data augmentation
test_transform = transforms.Compose([
# transforms.ToPILImage(), # 这个不需要
transforms.ToTensor(),
])
class ImgDataset(Dataset):
def __init__(self, x, y=None, transform=None):
self.x = x
# label is required to be a LongTensor 这里不知道为啥了
self.y = y # 本来是uint8
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
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), # 对应的out_channels
nn.ReLU(), # 非线性变换，一定要有
nn.MaxPool2d(2, 2, 0), # [64, 64, 64]
nn.Conv2d(64, 128, 3, 1, 1), # [128, 64, 64]
nn.BatchNorm2d(128),
nn.ReLU(),
nn.MaxPool2d(2, 2, 0), # [128, 32, 32]
nn.Conv2d(128, 256, 3, 1, 1), # [256, 32, 32]
nn.BatchNorm2d(256),
nn.ReLU(),
nn.MaxPool2d(2, 2, 0), # [256, 16, 16]
nn.Conv2d(256, 512, 3, 1, 1), # [512, 16, 16]
nn.BatchNorm2d(512),
nn.ReLU(),
nn.MaxPool2d(2, 2, 0), # [512, 8, 8]
nn.Conv2d(512, 512, 3, 1, 1), # [512, 8, 8]
nn.BatchNorm2d(512),
nn.ReLU(),
nn.MaxPool2d(2, 2, 0), # [512, 4, 4]
)
self.fc = nn.Sequential(
nn.Linear(512 * 4 * 4, 1024),
nn.ReLU(),
nn.Linear(1024, 512),
nn.ReLU(),
nn.Linear(512, 11)
)
def forward(self, x):
out = self.cnn(x) # 先过cnn
out = out.view(out.size()[0], -1) # out.size()[0] batch_size,拉成一列，过fc，比如说原本有32个，resize(8, -1) 8 * 4
return self.fc(out)
workspace_dir = './test'
train_x, train_y = readfile(os.path.join(workspace_dir, "training"), True) # 这里拼接字符串
val_x, val_y = readfile(os.path.join(workspace_dir, "validation"), True)
test_x = readfile(os.path.join(workspace_dir, "testing"), False)
batch_size = 10 # 这个设置太大会报错
train_set = ImgDataset(train_x, train_y, train_transform) # 扔进类中做数据transform
val_set = ImgDataset(val_x, val_y, test_transform)
train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True) # 返回一个可迭代对象，每次取出指定数目的数据，下次是否随机取
val_loader = DataLoader(val_set, batch_size=batch_size, shuffle=False) # 验证集就不随机了
model = Classifier().cuda() # 调用GPU
# print(model.load_state_dict(torch.load('cnn_params.pkl'))) # 载入已经训练好的模型
loss = nn.CrossEntropyLoss() # 调用现成的classification task，所以 loss 使用 CrossEntropyLoss
optimizer = torch.optim.Adam(model.parameters(), lr=0.001) # optimizer 使用 Adam，超参数
num_epoch = 10
for epoch in range(num_epoch):
epoch_start_time = time.time()
train_acc = 0.0
train_loss = 0.0
val_acc = 0.0
val_loss = 0.0
model.train() # 固定的写法，训练时候开启训练模式。
for i, data in enumerate(train_loader):
# data list data[0] 10 * 3 * 128 * 128 data[1](标签) 10,
optimizer.zero_grad() # 每次进行梯度归0，否者就会累加梯度
train_pred = model(data[0].cuda()) # 数据也用GPU，将数据丢给模型，自动调用forward函数
batch_loss = loss(train_pred, data[1].cuda()) # 計算 loss （注意 prediction 跟 label 必須同時在 CPU 或是 GPU 上） 10 * 11
batch_loss.backward() # 固定写法 利用 back propagation 算出每個參數的 gradient
optimizer.step() # 固定写法 以 optimizer 用 gradient 更新參數值
train_acc += np.sum(np.argmax(train_pred.cpu().data.numpy(), axis=1) == data[1].numpy()) # argmx找出每一行最大值对应的index
train_loss += batch_loss.item() # 取batch_loss的数值
model.eval() # 固定的写法，验证时候开启评估模式。
with torch.no_grad(): # 不计算梯度了，
for i, data in enumerate(val_loader):
val_pred = model(data[0].cuda())
batch_loss = loss(val_pred, data[1].cuda())
val_acc += np.sum(np.argmax(val_pred.cpu().data.numpy(), axis=1) == data[1].numpy())
val_loss += batch_loss.item()
# 將結果 print 出來
print('[%03d/%03d] %2.2f sec(s) Train Acc: %3.6f Loss: %3.6f | Val Acc: %3.6f loss: %3.6f' % \
(epoch + 1, num_epoch, time.time() - epoch_start_time, \
train_acc / train_set.__len__(), train_loss / train_set.__len__(), val_acc / val_set.__len__(),
val_loss / val_set.__len__())) # 打印一堆数据。
# torch.save(model.state_dict(), 'cnn_params.pkl') # 把训练好的模型保存起来
# 训练数据和验证数据拼在一起。再次训练已经训练好的模型（更新他的参数）。
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)
model_best = Classifier().cuda()
loss = nn.CrossEntropyLoss() # 因為是 classification task，所以 loss 使用 CrossEntropyLoss
optimizer = torch.optim.Adam(model_best.parameters(), lr=0.001) # optimizer 使用 Adam
num_epoch = 10
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
