【写给自己】记录一下第一次搭建简单CNN的过程

【写给自己】记录一下第一次搭建简单CNN的过程

用CNN实现分类

数据预处理

由于还没有拿到数据，所以先拿mnist数据集来预先练习一下。首先用代码将mnist处理成“one folder one class”的形式：

对于这种形式的分类，pytorch的ImageFolder可以完美解决数据预处理。只需要

from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader
path = './mnist'
transform = ...
train = ImageFolder(path, transform=transform)
train = Dataloader(train, batch_size=4, shuffle=True) #我的3070居然batch_size不能太大

这样子我们就得到了预处理好的数据。接下来就是配置网络了。

网络训练

class Model(nn.Module):
    def __init__(self, num_labels=10):
        super().__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        self.layer2 = nn.Sequential(
            nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        self.fc = nn.Linear(7*7*32, num_labels)

    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = out.reshape(out.size(0), -1)  # flatten或者这个方法来扁平化数据
        out = self.fc(out)
        return out
        
def fit(epoch, model, trainloader, device):
    correct = 0
    total = 0
    running_loss = 0

    model.train()  # 训练阶段
    for x, y in trainloader:
        x, y = x.to(device), y.to(device)
        y_pred = model(x)
        loss = loss_func(y_pred, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        with torch.no_grad():
            y_pred = torch.argmax(y_pred, dim=1)
            correct += (y_pred == y).sum().item()
            total += y.size(0)
            running_loss += loss.item()

    epoch_acc = correct / total
    epoch_loss = running_loss / len(trainloader.dataset)

    print('epoch: ', epoch + 1,
          'loss: ', round(epoch_loss, 3),
          'accuracy: ', round(epoch_acc, 3))

    return epoch_loss, epoch_acc

loss_func = torch.nn.CrossEntropyLoss()
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    model = Model()
    model.to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    epochs = 5
    train_loss = []
    train_acc = []

    for epoch in range(epochs):
        epoch_loss, epoch_acc= fit(epoch, model, train, device)
        train_loss.append(epoch_loss)
        train_acc.append(epoch_acc)

训练好之后就可以保存了

# 这个是保存全部参数的
# 建议是保存部分
torch.save(model, 'cnn.pth') # 不能带中文路径

训练部分完整代码

import matplotlib.pyplot as plt
import numpy as np
import torch.nn as nn
import torch
from torch.utils.data import DataLoader,Dataset
from torchvision import transforms, utils
from torchvision.datasets import ImageFolder
import time

def one_hot(labels):
    pass

# all images in one folder & each class one folder用ImageFolder
class Model(nn.Module):
    def __init__(self, num_labels=10):
        super().__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        self.layer2 = nn.Sequential(
            nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        self.fc = nn.Linear(7*7*32, num_labels)

    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = out.reshape(out.size(0), -1)  # flatten或者这个方法来扁平化数据
        out = self.fc(out)
        return out


def fit(epoch, model, trainloader, device):
    correct = 0
    total = 0
    running_loss = 0

    model.train()  # 训练阶段
    for x, y in trainloader:
        x, y = x.to(device), y.to(device)
        y_pred = model(x)
        loss = loss_func(y_pred, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        with torch.no_grad():
            y_pred = torch.argmax(y_pred, dim=1)
            correct += (y_pred == y).sum().item()
            total += y.size(0)
            running_loss += loss.item()

    epoch_acc = correct / total
    epoch_loss = running_loss / len(trainloader.dataset)

    print('epoch: ', epoch + 1,
          'loss: ', round(epoch_loss, 3),
          'accuracy: ', round(epoch_acc, 3))

    return epoch_loss, epoch_acc


if __name__ == "__main__":

    path = './mnist_test'

    transform=transforms.Compose([
        transforms.ToTensor(),  # 将图片转换为Tensor,归一化至[0,1]
        transforms.Normalize(mean=[.5, .5, .5], std=[.5, .5, .5])
    ])

    train = ImageFolder(path, transform=transform)
    train = DataLoader(train, batch_size=4, shuffle=True)

    loss_func = torch.nn.CrossEntropyLoss()
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    model = Model()
    model.to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    epochs = 5
    train_loss = []
    train_acc = []

    for epoch in range(epochs):
        epoch_loss, epoch_acc= fit(epoch, model, train, device)
        train_loss.append(epoch_loss)
        train_acc.append(epoch_acc)

    print("done")
    save_path = "\CNN_work"  # 不能带中文
    torch.save(model, 'cnn.pth')

调用已经训练好的模型

首先需要加载训练好的网络

test_model = torch.load('cnn.pth', map_location='cuda')

然后就可以加载测试数据，来判断类别了，这里我还是用了一样的数据集。

test_data = cv2.imread('./mnist_test/5/23.png')

trans = transforms.ToTensor()
img = trans(test_data)
img = img.to('cuda')
img = img.view(1, 3, 28, 28)	# 把数据变成对应尺寸

print('done')
out = test_model(img)
out = F.softmax(out)
proba, class_ind = torch.max(out, 1)

proba = float(proba)
class_ind = int(class_ind)
print(proba, class_ind)

结果如图

还需要学会自定义Dataset的方式

这个up写的很好很详细https://www.bilibili.com/video/BV1354y1s7kQ