pytorch（一）

Pytorch模型训练

---

前言

pytorch的简易训练模型方法，适用于大多数图片训练，导入数据集即可，前提是数据集是图片，程序中只需要添加相应的训练集路径和测试集路径，并改变class中的类别，即可实现其功能，其余不需要改变。

---

Pytorch程序

话不多说，直接上程序
代码如下：

import time
import torch
import torchvision.datasets
from PIL import Image
from torch import nn
from torch.optim import optimizer
from torch.utils.data import DataLoader, Dataset
from torch.utils.tensorboard import SummaryWriter

# 加载数据集
def loadtraindata():
    path = r"data/train"  # 训练数据集路径
    trainset = torchvision.datasets.ImageFolder(path,
                                                transform=torchvision.transforms.Compose([
                                                    torchvision.transforms.Resize((32, 32)),
                                                    torchvision.transforms.CenterCrop(32),
                                                    torchvision.transforms.ToTensor()])
                                                )

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

def loadtestdata():
    path1 = r"data/valid" # 测试数据集路径
    testset = torchvision.datasets.ImageFolder(path1,
                                                transform=torchvision.transforms.Compose([
                                                    torchvision.transforms.Resize((32, 32)),
                                                    torchvision.transforms.CenterCrop(32),
                                                    torchvision.transforms.ToTensor()])
                                                )

    testloader = torch.utils.data.DataLoader(testset, batch_size=4,
                                              shuffle=True, num_workers=2)
    return testloader

# 搭建神经网络
class shenjing(nn.Module) :
    def __init__(self):
        super(shenjing, self).__init__()
        self.model = nn.Sequential(nn.Conv2d(3, 32, 5, padding=2),
                                   nn.MaxPool2d(2),
                                   nn.Conv2d(32, 32, 5, padding=2),
                                   nn.MaxPool2d(2),
                                   nn.Conv2d(32, 64, 5, padding=2),
                                   nn.MaxPool2d(2),
                                   nn.Flatten(),
                                   nn.Linear(1024, 64),
                                   nn.Linear(64, 4)
                                  )

    def forward(self ,x):
        x = self.model(x)
        return x

if __name__ == '__main__' :

    shendu = shenjing()
    shendu = shendu.cuda()

    classes = ('APPle_scab','Black_rot', 'Cedar_apple_rust', 'healthy')#我这里是4组，根据自己的判断进行输入

    #计算数据集的长度
    train_data_size = len(loadtraindata())
    test_data_size = len(loadtestdata())
    print("训练数据集的长度为：{}".format(train_data_size))
    print("测试数据集的长度为: {}".format(test_data_size))

    #损失函数
    loss_fn = nn.CrossEntropyLoss()
    loss_fn = loss_fn.cuda()

    #优化器
    learn_rate = 1e-3
    optimizer = torch.optim.Adam(shendu.parameters() ,lr=learn_rate)

    # 记录的次数
    total_train_step = 0
    total_test_step = 0
    epoch = 100

    # 添加tensorboard
    # writer = SummaryWriter("moxing1")
    # start_time = time.time()
    # 训练开始
    shendu.train()
    for i in range(epoch) :
        print("---------第{}轮训练开始---------".format(i + 1))
        for data in loadtraindata() :
            imgs,targets = data
            imgs = imgs.cuda()
            targets = targets.cuda()
            output = shendu(imgs)
            loss = loss_fn(output,targets)
            #优化器优化模型
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            total_train_step = total_train_step + 1
            # if total_train_step % 100 == 0:
            #     end_time = time.time()
            #     print(end_time - start_time)
            #     print("训练次数:{}，损失误差:{}".format(total_train_step ,loss.item()))
            #     writer.add_scalar("train_loss",loss.item(),total_train_step)

        # 11、测试模型
        shendu.eval()
        total_test_loss = 0
        total_accuracy = 0
        with torch.no_grad() :
            for data in loadtestdata() :
                imgs,targets = data
                imgs = imgs.cuda()
                targets = targets.cuda()
                outputs = shendu(imgs)
                loss = loss_fn(outputs,targets)
                total_test_loss = total_test_loss + loss.item()
                accuracy = (outputs.argmax(1) == targets).sum()
                total_accuracy = total_accuracy + accuracy
        # print("整体测试集上的损失误差:{}".format(total_test_loss))
        print("整体测试集上的正确率:{}%".format(100*total_accuracy/4/test_data_size))
        # writer.add_scalar("test_loss",total_test_loss,total_test_step)
        # writer.add_scalar("test_accuray",total_accuracy ,total_test_step)
        total_test_step = total_test_step + 1

        torch.save(shendu , "shendu_{}.pth".format(i + 1))
        #torch.save(shendu.state_dict(),"shendu_{}.pth".format（i+ 1）)
        print("模型已保存")

    # writer.close()