pytorch 10分类GPU训练 完整代码

import os.path
from typing import Iterator
import numpy as np
import torch
import cv2
from PIL import Image
from torch.utils.data import Dataset,DataLoader,Subset,random_split
import re
from functools import reduce
from torch.utils.tensorboard import SummaryWriter as Writer
from torchvision import transforms,datasets
import torchvision as tv
from torch import nn
import torch.nn.functional as F
import time
t1=transforms.ToTensor()
t2=transforms.Normalize([0.5,0.5,0.5],[0.5,0.5,0.5])
t=transforms.Compose([t1,t2])

#数据集下载后为PIL格式：
TrainDataSet=datasets.CIFAR10(root="./myDataSet",train=True,download=True,transform=t)
#取最后320个作为验证集方式1：取子集
#validDataSet=Subset(TrainDataSet,[x for x in range(len(TrainDataSet))][-320:])
#或直接使用随机划分的方法
TrainDataSet,ValidDataSet=random_split(TrainDataSet,[int(len(TrainDataSet)*0.9),len(TrainDataSet)-int(len(TrainDataSet)*0.9)])
TextDataSet=datasets.CIFAR10(root="./myDataSet",train=False,download=True,transform=t)
dataloaderForText=DataLoader(TextDataSet,shuffle=True,batch_size=32,num_workers=0,drop_last=False)
dataloaderForTrain=DataLoader(TrainDataSet,shuffle=True,batch_size=32,num_workers=0,drop_last=False)
dataloaderForValid=DataLoader(ValidDataSet,shuffle=True,batch_size=32,num_workers=0,drop_last=False)

#查看命令：tensorboard --logdir=./myBorderText
#可用pycharm中code中的generater功能实现：
class myCustomerNetWork(nn.Module):
    def __init__(self):
        super().__init__()
        #输入3通道输出6通道：
        self.features=nn.Sequential(nn.Conv2d(3, 64, (3, 3)),nn.ReLU(),nn.Conv2d(64,128,(3,3)),
                                    nn.ReLU(),nn.Conv2d(128,256,(3,3)),nn.ReLU(),nn.AdaptiveAvgPool2d(1))

        self.classfired=nn.Sequential(nn.Flatten(),nn.Linear(256,80),nn.Dropout(),nn.Linear(80,10))
    def forward(self,x):
        return self.classfired(self.features(x))

myNet=myCustomerNetWork()
lossFun = nn.CrossEntropyLoss()
#模型转换至GPU上进行运算：需要转换的对象：模型，数据,损失函数
if torch.cuda.is_available():
    myNet=myNet.cuda()
    lossFun=lossFun.cuda()
# #载入日志写入器：
writer=Writer("./myBorderText")
optim=torch.optim.Adam(myNet.parameters(),lr=0.001)
for epoch in range(100):
    lossOneEpoch=0
    myNet.train()
    rightNumAll=0
    for index,(data,label) in enumerate(dataloaderForTrain):
        if torch.cuda.is_available():
            data=data.cuda()
            label = label.cuda()
        outPut=myNet(data)
        loss = lossFun(outPut, label.long())
        optim.zero_grad()
        loss.backward()
        optim.step()
        lossOneEpoch+=loss
        rightNumAll += (outPut.argmax(1) == label).sum().item()
        accuracy = rightNumAll * 1.0 / len(TrainDataSet)
        print('\r目前训练至第%d代%d/%d批' % (epoch,index,dataloaderForTrain.__len__()), end="")
    print()
    #当一次显示一个数据集时，使用add_images方法：
    writer.add_scalar("训练集损失函数值", scalar_value=lossOneEpoch.float(), global_step=epoch)
    writer.add_scalar("训练集预测精度值", scalar_value=accuracy, global_step=epoch)
    print("训练集的预测精度值为{}，损失值为{}".format(accuracy, lossOneEpoch.float()))
    #切换为测试模式
    myNet.eval()
    rightNumAll=0
    lossOneEpoch=0
    #不再更新梯度
    with torch.no_grad():
        for data,label in dataloaderForText:
            if torch.cuda.is_available():
                data = data.cuda()
                label = label.cuda()
            outPut=myNet(data)
            loss = lossFun(outPut, label.long())
            lossOneEpoch += loss
            #返回是布尔值Tensor,将其中为true的相加
            rightNumAll+=(outPut.argmax(1)==label).sum().item()
        accuracy= rightNumAll*1.0/len(TextDataSet)
    writer.add_scalar("测试集精确度值", scalar_value=accuracy, global_step=epoch)
    writer.add_scalar("测试集损失函数值", scalar_value=lossOneEpoch.float(), global_step=epoch)
    print("测试集的预测精度值为{}，损失值为{}".format(accuracy,lossOneEpoch.float()))
    #当训练精度大于0.9时停止训练
    if accuracy>0.9:
        torch.save(myNet.state_dict(),"./saveTextOnlyParams.pth")
        break
writer.close()
#在验证集上进行测试：
myNet.eval()
rightNumAll=0
#不再更新梯度
with torch.no_grad():
    for data,label in dataloaderForValid:
        if torch.cuda.is_available():
            data=data.cuda()
            label=label.cuda()
        outPut=myNet(data)
        loss = lossFun(outPut, label.long())
        #返回是布尔值Tensor,将其中为true的相加
        rightNumAll+=(outPut.argmax(1)==label).sum().item()
    accuracy= rightNumAll*1.0/len(ValidDataSet)
print("最终精度为{}".format(accuracy))