pytorch实现手写字的代码

小咪咪 2023-03-12 15:25 10阅读 0赞

## 一、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)}")

测试后的输出结果为：  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzM5ODUyNjc2_size_16_color_FFFFFF_t_70]

## 二、代码重构 ##

从上述的代码上可以大致看出并没有对网络结构进行一个很好的处理，而是直接采用`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)

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

[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzM5ODUyNjc2_size_16_color_FFFFFF_t_70]: /images/20230312/3cba4bc929344c34a837918f1a33f330.png