训练、测试、验证

同一个目录下
model文件

#搭建神经网络
import torch
from torch import nn


class Tudui(nn.Module):
    def __init__(self):
        super(Tudui, self).__init__()
        self.Module =nn.Sequential(
            nn.Conv2d(3,32,5,1,padding=2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 32, 5, 1, padding=2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 5, 1, padding=2),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(64*4*4,64),
            nn.Linear(64, 64)

        )
    def forward(self,x):
        x = self.Module(x)
        return x
tudui = Tudui()
input = torch.ones(64,3,32,32)    #64张图片，选取3通道，大小是32*32
print(input)
output = tudui(input)
print(output)

train_GPU文件

import torch
from torch.utils.tensorboard import SummaryWriter

from model import *   #神经网络的设计写在同文件夹下的model文件

#准备数据集
import torchvision
from torch import nn
from torch.nn import Conv2d,MaxPool2d,Flatten,Linear
from torch.utils.data import DataLoader

import  time

# data_train = torchvision.datasets.CIFAR10("./dataset",train=True,transform=torchvision.transforms.ToTensor(),download=True)   #数据集没有下载的话
data_train = torchvision.datasets.CIFAR10("./dataset",train=True,transform=torchvision.transforms.ToTensor())
data_test = torchvision.datasets.CIFAR10("./dataset",train=False,transform=torchvision.transforms.ToTensor())
print("训练集长度{},测试集长度{}".format(len(data_train),len(data_test)))

#dataloader加载数据集
train_dataloader = DataLoader(data_train,batch_size=64)
test_dataloader = DataLoader(data_test,batch_size=64)

#GPU加速，{模型，数据（输入、标签），损失函数}.to（device）
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

#创建网络模型
tudui = Tudui()
tudui = tudui.to(device)
print(tudui)

#创建损失函数   交叉熵损失函数
loss_fn = nn.CrossEntropyLoss()
loss_fn = loss_fn.to(device)

#创建优化器   随机下降
optimizer = torch.optim.SGD(tudui.parameters(),lr=0.01)

##设置训练网络的一些参数
#记录训练的次数
total_train_step=0
#记录测试的次数
total_test_step=0
#训练的轮数
epoch = 10


#添加tensorboard
writer = SummaryWriter("tensorboard")


start_time = time.time()
for i in range(epoch):
    print("--------------------------训练的第{}几轮---------------------------".format(i+1))

    #训练步骤开始
    for data in train_dataloader:
        imgs,targets = data
        imgs = imgs.to(device)
        targets = targets.to(device)
        outputs = tudui(imgs)
        loss = loss_fn(outputs,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("训练次数：{},loss:{},结束时间：{}".format(total_train_step,loss,end_time-start_time))
            writer.add_scalar("train_loss",loss,total_train_step)



    #测试步骤开始
    total_test_loss=0
    total_accuracy = 0

    with torch.no_grad():   #清零
        for data in test_dataloader:
            imgs,targets = data
            imgs = imgs.to(device)
            targets = targets.to(device)

            outputs = tudui(imgs)
            loss = loss_fn(outputs,targets)
            total_test_loss = total_test_loss+loss

            accuracy = (outputs.argmax(1) == targets).sum()
            total_accuracy = total_accuracy + accuracy


        print("整体测试集上的loss:{}".format(total_test_loss))
        print("整体测试集上的accuracy:{}".format(total_accuracy))

        writer.add_scalar("test_loss", total_test_loss, i+1)
        writer.add_scalar("test_accuracy", total_accuracy, i+1)


        torch.save(tudui,"tudui_{}.pth".format(i+1))

writer.close()  #cmd输入    tensorboard --logdir=logs   logs改成绝对路径    E:\pythonProject2\tensorboard

test文件

import torch
import torchvision.transforms
from PIL import Image
import cv2

img_path = "dataset/dog.png"   #  ../是退回上一级的
# print(img_path)
# img = cv2.imread("img_path",cv2.IMREAD_COLOR)
img = Image.open('dataset/dog.png')

transform = torchvision.transforms.Compose([
    torchvision.transforms.Resize((32,32)),
    torchvision.transforms.ToTensor()
])
image = transform(img)



device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

model = torch.load("tudui_8.pth").to(device)
image = torch.reshape(image,[1,3,32,32]).to(device)   #1张图片，用的3通道，大小是32*32

model.eval()
with torch.no_grad():
    output = model(image)
print(output)