pytorch实现手写字的代码

一、Pytorch手写字代码

废话不多说，直接上代码。

import torch
from torchvision import datasets, transforms
from torch import nn
from torch import optim
import numpy as np

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

dataset = datasets.MNIST("minist_data", download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(dataset, batch_size=6)

model = nn.Sequential(nn.Linear(784, 256),
                      nn.ReLU(),
                      nn.Linear(256, 64),
                      nn.ReLU(),
                      nn.Linear(64, 10),
                      nn.LogSoftmax(dim=1))

criterion = nn.NLLLoss()
optimizer = optim.SGD(model.parameters(), lr=0.005)
epoch = 5
for e in range(epoch):
    running_loss = 0
    for images, labels in trainloader:
        images = images.view(images.shape[0], -1)
        optimizer.zero_grad()
        output = model.forward(images)
        loss = criterion(output, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
    else:
        print(f"第{e}代, 训练损失：{running_loss/len(trainloader)}")

测试后的输出结果为：

二、代码重构

从上述的代码上可以大致看出并没有对网络结构进行一个很好的处理，而是直接采用nn.Sequential按照顺序放入网络，下边我对代码进行重构，使其上升成为一个比较规范的项目，以便在基础上对网络进行修改、参数修改，以及其他代码的补充等，重构后的代码如下。

# coding=utf-8

import torch
from torchvision import datasets, transforms
import torch.nn as nn
from torch import optim
from torch.utils.data.dataloader import default_collate


class Models(torch.nn.Module):

    def __init__(self):
        super(Models, self).__init__()
        self.connect1 = nn.Linear(784, 256)
        self.connect2 = nn.Linear(256, 64)
        self.connect3 = nn.Linear(64, 10)
        self.softmax = nn.LogSoftmax(dim=1)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.connect1(x)
        x = self.relu(x)
        x = self.connect2(x)
        x = self.relu(x)
        x = self.connect3(x)
        x = self.softmax(x)
        return x


class Test:

    def __init__(self):
        self.epoch = 5
        self.batch_size = 6
        self.learning_rate = 0.005
        self.models = Models()

    def transdata(self):
        transform = transforms.Compose(
            [transforms.ToTensor(), transforms.Normalize((0.5, ), (0.5, ))])
        return transform

    def loaddata(self):
        dataset = datasets.MNIST(
            "minist_data", download=True, transform=self.transdata())
        dataset = torch.utils.data.DataLoader(
            dataset, batch_size=self.batch_size)
        return dataset

    def lossfunction(self):
        criterion = nn.NLLLoss()
        return criterion


def main(datahandle, models):
    dataset = datahandle().loaddata()
    model = models()

    criterion = datahandle().lossfunction()
    optimizer = optim.SGD(model.parameters(), datahandle().learning_rate)
    epoch = datahandle().epoch

    for single_epoch in range(epoch):
        running_loss = 0
        for image, lable in dataset:
            image = image.view(image.shape[0], -1)

            optimizer.zero_grad()
            output = model(image)
            loss = criterion(output, lable)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()

        print(f"第{single_epoch}代, 训练损失：{running_loss/len(dataset)}")


if __name__ == '__main__':
    main(Test, Models)

三、保存模型，并测试

1、保存模型

在上述重构后的代码中仍旧没有加入模型的保存。在测试程序中加入对模型的保存和测试。
修改部分是在上述重构代码中的main函数中。

def main(datahandle, models):
    dataset = datahandle().loaddata()
    model = models()

    criterion = datahandle().lossfunction()
    optimizer = optim.SGD(model.parameters(), datahandle().learning_rate)
    epoch = datahandle().epoch

    for single_epoch in range(epoch):
        running_loss = 0
        for image, lable in dataset:
            image = image.view(image.shape[0], -1)

            optimizer.zero_grad()
            output = model(image)
            loss = criterion(output, lable)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()

        avg_running_loss = running_loss / len(dataset)
        torch.save(model, "model" + str(avg_running_loss) + ".pkl")
        print(f"第{single_epoch}代, 训练损失：{running_loss/len(dataset)}")

其中 torch.save(model, "model" + str(avg_running_loss) + ".pkl")会在当前文件夹下保存每个epoch的生成的权重文件。本次中保存是使用的是torch.save()函数，可以直接将网络和训练参数同时保存到pkl文件中，加载时只需通过torch.load()进行读取加载模型即可。

2、测试代码

class TestModel:
    """docstring for TestModel"""

    def __init__(self):
        self.image_normallization_mean = [0.5, ]
        self.image_normallization_std = [0.5, ]
        self.model_path = "/media/zhaokaiyue/_dde_data1/视频教学/代码/model0.09349177049578038.pkl"

    def init_transform_image(self):
        normalize = transforms.Normalize(
            self.image_normallization_mean, self.image_normallization_std)
        image_transform = transforms.Compose(
            [transforms.ToTensor(), normalize])
        return image_transform

    def init_model(self):
        if not os.path.exists(self.model_path):
            print("model is not exited!")
            return

        if not os.path.isfile(self.model_path):
            print("file is not file!")

        # device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        model = torch.load(self.model_path, map_location=“cpu”)
        model.eval()
        return model


def test_model(filedir):
    handlemodels = TestModel()
    handle = handlemodels.init_transform_image()
    model = handlemodels.init_model()

    with torch.no_grad():

        for img in os.listdir(filedir):
            child_file = os.path.join(filedir, img)
            img = handle(child_file)
            img = torch.autograd.variable(img)
            predict = model(img)
            print(predict)

上述测试代码是测试的整个文件夹下的图片，可以根据实际需求对代码进行修改，使满足适合自己的项目需求。