06 Pytorch代码实践之【LeNet手写数字识别】

LeNet:

训练部分，包括保存训练结果

Minst数据集下载（pytorch自带的可能比较慢）：http://yann.lecun.com/exdb/mnist/

import os
from torchvision import datasets, transforms
import torch
from torch import nn, optimM
from time import perf_counter

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(0.1307, 0.3081)
])
trainset = datasets.MNIST(
    'mnist', train=True, download=True, transform=transform)
testset = datasets.MNIST('mnist', train=False,
                         download=True, transform=transform)


class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.c1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=(5, 5))
        self.c3 = nn.Conv2d(in_channels=6, out_channels=16, kernel_size=(5, 5))
        self.fc1 = nn.Linear(in_features=16 * 4 * 4, out_features=120)
        self.fc2 = nn.Linear(in_features=120, out_features=84)
        self.fc3 = nn.Linear(in_features=84, out_features=10)

    def forward(self, x):
        x = nn.functional.max_pool2d(nn.functional.relu(self.c1(x)), 2)
        x = nn.functional.max_pool2d(
            nn.functional.relu(self.c3(x)), 2)  # [b,16,4,4]
        # 按照图片分别展平[b,16,4,4]->[b,256],b代表b张图片
        x = x.view(-1, self.num_flat_features(x))
        x = nn.functional.relu(self.fc1(x))  # [b,256]->[b,120]
        x = nn.functional.relu(self.fc2(x))  # [b,120]->[b,84]
        x = self.fc3(x)  # [b,84]->[b,10]
        return x

    def num_flat_features(self, x):
        size = x.size()[1:]
        num_feature = 1
        for s in size:
            num_feature *= s
        return num_feature


CUDA = torch.cuda.is_available()
if CUDA:
    lenet = LeNet().cuda()
else:
    lenet = LeNet()

optimizer = optim.SGD(lenet.parameters(), lr=0.001, momentum=0.9)
criterion = nn.CrossEntropyLoss()

trainloader = torch.utils.data.DataLoader(trainset, batch_size=4, shuffle=True)


# 训练
def train(model, criterion, optimizer, epochs=1):
    for epoch in range(epochs):
        running_loss = 0.0
        for i, data in enumerate(trainloader, 0):
            inputs, labels = data
            if CUDA:
                inputs, labels = inputs.cuda(), labels.cuda()
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()
            if i % 1000 == 999:
                print('[Epoch:%d,Batch:%5d] Loss:%.3f' %
                      (epoch + 1, i + 1, running_loss / 1000))
                running_loss = 0.0
    print("Fininsh Training!")

    return model, loss


start = perf_counter()
model, loss = train(lenet, criterion, optimizer, epochs=2)
finish = perf_counter()
time = finish - start
print("计算时间：%s" % time)
print("final loss:", loss.item())
# print("weights:", list(model.parameters()))

# ---------------------------存储---------------------------


# 加载模型
def load_param(model, path):
    if os.path.exists(path):
        model.load_state_dict(torch.load(path))


# 保存模型
def save_param(model, path):
    torch.save(model.state_dict(), path)


save_param(lenet, 'model.pkl')

输出：

[Epoch:1,Batch: 1000] Loss:1.312
[Epoch:1,Batch: 2000] Loss:0.331
[Epoch:1,Batch: 3000] Loss:0.215
[Epoch:1,Batch: 4000] Loss:0.172
[Epoch:1,Batch: 5000] Loss:0.113
[Epoch:1,Batch: 6000] Loss:0.130
[Epoch:1,Batch: 7000] Loss:0.115
[Epoch:1,Batch: 8000] Loss:0.117
[Epoch:1,Batch: 9000] Loss:0.115
[Epoch:1,Batch:10000] Loss:0.088
[Epoch:1,Batch:11000] Loss:0.106
[Epoch:1,Batch:12000] Loss:0.082
[Epoch:1,Batch:13000] Loss:0.068
[Epoch:1,Batch:14000] Loss:0.081
[Epoch:1,Batch:15000] Loss:0.081
[Epoch:2,Batch: 1000] Loss:0.073
[Epoch:2,Batch: 2000] Loss:0.052
[Epoch:2,Batch: 3000] Loss:0.066
[Epoch:2,Batch: 4000] Loss:0.068
[Epoch:2,Batch: 5000] Loss:0.065
[Epoch:2,Batch: 6000] Loss:0.061
[Epoch:2,Batch: 7000] Loss:0.062
[Epoch:2,Batch: 8000] Loss:0.062
[Epoch:2,Batch: 9000] Loss:0.063
[Epoch:2,Batch:10000] Loss:0.055
[Epoch:2,Batch:11000] Loss:0.053
[Epoch:2,Batch:12000] Loss:0.058
[Epoch:2,Batch:13000] Loss:0.054
[Epoch:2,Batch:14000] Loss:0.057
[Epoch:2,Batch:15000] Loss:0.055
Fininsh Training!
计算时间：142.44309929999872
final loss: 0.0002086162567138672

测试部分

import os
import torch
from torchvision import datasets, transforms
from torch import nn, optim

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(0.1307, 0.3081)
])


class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.c1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=(5, 5))
        self.c3 = nn.Conv2d(in_channels=6, out_channels=16, kernel_size=(5, 5))
        self.fc1 = nn.Linear(in_features=16 * 4 * 4, out_features=120)
        self.fc2 = nn.Linear(in_features=120, out_features=84)
        self.fc3 = nn.Linear(in_features=84, out_features=10)

    def forward(self, x):
        x = nn.functional.max_pool2d(nn.functional.relu(self.c1(x)), 2)
        x = nn.functional.max_pool2d(nn.functional.relu(self.c3(x)), 2)  # [b,16,4,4]
        x = x.view(-1, self.num_flat_features(x))  # 按照图片分别展平[b,16,4,4]->[b,256],b代表b张图片
        x = nn.functional.relu(self.fc1(x))  # [b,256]->[b,120]
        x = nn.functional.relu(self.fc2(x))  # [b,120]->[b,84]
        x = self.fc3(x)  # [b,84]->[b,10]
        return x

    def num_flat_features(self, x):
        size = x.size()[1:]
        num_feature = 1
        for s in size:
            num_feature *= s
        return num_feature


CUDA = torch.cuda.is_available()
if CUDA:
    lenet = LeNet().cuda()
else:
    lenet = LeNet()

optimizer = optim.SGD(lenet.parameters(), lr=0.001, momentum=0.9)
criterion = nn.CrossEntropyLoss()


# 加载模型
def load_param(model, path):
    if os.path.exists(path):
        model.load_state_dict(torch.load(path))


# 保存模型
def save_param(model, path):
    torch.save(model.state_dict(), path)


# ---------------------------测试---------------------------
testset = datasets.MNIST('mnist', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4, shuffle=False)


def test(testloader, model):
    correct = 0
    total = 0
    for data in testloader:
        images, labels = data
        if CUDA:
            images = images.cuda()
            labels = labels.cuda()
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum()
        print("Accuracy on the test set :%d %%" % (100 * correct / total))


load_param(lenet, 'model.pkl')
test(testloader, lenet)

..\aten\src\ATen\native\BinaryOps.cpp:81: UserWarning: Integer division of tensors using div or / is deprecated, and in a future release div will perform true division as in Python 3. Use true_divide or floor_divide (// in Python) instead.
Accuracy on the test set :100 %
Accuracy on the test set :100 %
Accuracy on the test set :100 %
Accuracy on the test set :100 %
Accuracy on the test set :100 %
Accuracy on the test set :100 %
Accuracy on the test set :100 %
Accuracy on the test set :100 %
               # 省略很多。。。。。。。。。。。。。。。。。。
Accuracy on the test set :99 %
Accuracy on the test set :99 %
Accuracy on the test set :99 %
Accuracy on the test set :99 %
Accuracy on the test set :99 %
Accuracy on the test set :99 %
Accuracy on the test set :99 %
Accuracy on the test set :99 %
Accuracy on the test set :98 %

遇到的问题：

trainloader = torch.utils.data.DataLoader(trainset, batch_size=4, shuffle=True, num_workers=2)

中的num_workers=2表示同时启动多线程，此处为2，可能会出现以下问题：

RuntimeError: 
        An attempt has been made to start a new process before the
        current process has finished its bootstrapping phase.

        This probably means that you are not using fork to start your
        child processes and you have forgotten to use the proper idiom
        in the main module:

            if __name__ == '__main__':
                freeze_support()
                ...

        The "freeze_support()" line can be omitted if the program
        is not going to be frozen to produce an executable.

具体原因不清楚，但是关闭多线程，使用单一线程即可运行，即删除num_workers=2