参与拿奖：本文已参与「新人创作礼」活动，一起开启掘金创作之路
PS:完整代码文末自取

1.训练准备

1.1 导入相关包

import torch
from torch import nn
import torchvision
from torchvision import transforms
from torch.utils.tensorboard import SummaryWriter
from time import time

1.2 加载数据集

# 数据集
train_data = torchvision.datasets.CIFAR10('./data/05 data', train=True, transform=transforms.ToTensor(), download=True)
test_data = torchvision.datasets.CIFAR10('./data/05 data', train=False, transform=transforms.ToTensor(), download=True)

print("train_data's length: {}".format(len(train_data)))
print("test_data's length: {}".format(len(test_data)))

# 利用dataloader加载数据
from torch.utils.data import DataLoader

train_data = DataLoader(train_data
                        , batch_size=64  # 每次抽取数据的大小
                        , shuffle=False  # 是否打乱顺序
                        )
test_data = DataLoader(test_data
                       , batch_size=64
                       , shuffle=False)

1.3 搭建神经网络

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

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

1.4 参数初始化

# 网络实例化
mynet = myNet()

# 损失
loss_fn = nn.CrossEntropyLoss()

# 优化器
learning_rate = 0.01  # 等价于1e-2
optimizer = torch.optim.SGD(mynet.parameters(), lr=learning_rate)

# 设置训练网络参数
total_train_step = 0  # 训练与测试次数
total_test_step = 0

epoch = 3  # 训练轮数

1.5 训练

# 使用tensorboard绘制图像
writer = SummaryWriter("myNet")

# 训练步骤
for i in range(epoch):
    print("..................train's round %d............" % (i + 1))
    begin = time()
    mynet.train()  # 启动训练模式
    for data in train_data:
        imgs, targets = data
        outputs = mynet(imgs)
        loss = loss_fn(outputs, targets)  # 损失

        optimizer.zero_grad()  # 初始化梯度（调优）
        loss.backward()  # 反向传播
        optimizer.step()

        total_train_step += 1
        if total_train_step % 100 == 0:  # 每1训练00次打印一次
            print(time() - begin)
            print("train round {},Loss {}".format(total_train_step, loss))
            writer.add_scalar("train_loss", loss.item(), total_train_step)

1.6 测试

# 测试步骤
mynet.eval()  # 启动调试模式
total_test_loss = 0  # 测试的损失
total_accuracy = 0
with torch.no_grad():
    for data in test_data:
        imgs, targets = data
        outputs = mynet(imgs)
        loss = loss_fn(outputs, targets)  # 损失
        total_test_loss += loss.item()
        accuracy = (outputs.argmax(1) == targets).sum()  # 预判正确
        total_accuracy += accuracy
print('total test accuracy {}'.format(total_accuracy / len(test_data)))  # 测试正确率
print("total test loss {}".format(total_test_loss))
writer.add_scalar("test_loss", total_test_loss, total_test_step)
total_test_step += 1

# 保存每次的模型
torch.save(mynet, "./data/05 data/myNet/model_%d.pth" % i)  # 保存模型

2.使用GPU进行训练

2.1 训练准备

进入conda虚拟环境 ,在终端使用 nvidia-smi 查看cuda版本
使用.to(divice),torch.divice("cuda")

2.2 完整代码

# 可以通过 from model import* 导入上述模型
import torch
from torch import nn
import torchvision
from torchvision import transforms
from torch.utils.tensorboard import SummaryWriter
from time import time

print(torch.__version__,torch.cuda.is_available())


# 数据集
train_data=torchvision.datasets.CIFAR10('./data/05 data',train=True,transform=transforms.ToTensor(),download=True)
test_data=torchvision.datasets.CIFAR10('./data/05 data',train=False,transform=transforms.ToTensor(),download=True)

print("train_data's length: {}" .format(len(train_data)))
print("test_data's length: {}".format(len(test_data)))




# 利用dataloader加载数据
from torch.utils.data import DataLoader
train_data=DataLoader(train_data
                      ,batch_size=64 #每次抽取数据的大小
                      ,shuffle=False #是否打乱顺序
                     )
test_data=DataLoader(test_data
                     ,batch_size=64
                     ,shuffle=False)



# 搭建神经网络
class myNet(nn.Module):
    def __init__(self):
        super(myNet,self).__init__()
        self.model=nn.Sequential(
            nn.Conv2d(3,32,5,1,2)
            ,nn.MaxPool2d(2)
            ,nn.Conv2d(32,32,5,1,2)
            ,nn.MaxPool2d(2)
            ,nn.Conv2d(32,64,5,1,2)
            ,nn.MaxPool2d(2)
            ,nn.Flatten() # 64x4x4
            ,nn.Linear(64*4*4,64)
            ,nn.Linear(64,10)
        )
    def forward(self,x):
        output=self.model(x)
        return output
    

# 网络实例化
mynet=myNet()
if torch.cuda.is_available():
    mynet=mynet.cuda() #将模型转入cuda



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



# 优化器
learning_rate=0.01 #等价于1e-2
optimizer=torch.optim.SGD(mynet.parameters(),lr=learning_rate)

# 设置训练网络参数
total_train_step=0 # 训练与测试次数
total_test_step=0

epoch=3 # 训练轮数



# 使用tensorboard绘制图像
writer=SummaryWriter("myNet")

# 训练步骤
for i in range(epoch):
    print("..................train's round %d............"%(i+1))
    begin=time()
    mynet.train() #启动训练模式
    for data in train_data:
        imgs,targets=data
        if torch.cuda.is_available():
            imgs=imgs.cuda()
            targets=targets.cuda()
        outputs=mynet(imgs)
        loss=loss_fn(outputs,targets) #损失
        
        optimizer.zero_grad() #初始化梯度（调优）
        loss.backward() #反向传播
        optimizer.step()
        
        total_train_step+=1
        if total_train_step%100==0: #每1训练00次打印一次
            print(time()-begin)
            print("train round {},Loss {}".format(total_train_step,loss))
            writer.add_scalar("train_loss",loss.item(),total_train_step)
        
    
    # 测试步骤
    mynet.eval() # 启动调试模式
    total_test_loss=0 #测试的损失
    total_accuracy=0
    with torch.no_grad():
        for data in test_data:
            imgs,targets=data
            if torch.cuda.is_available():
                imgs=imgs.cuda()
                targets=targets.cuda()
            outputs=mynet(imgs)
            loss=loss_fn(outputs,targets) #损失
            total_test_loss+=loss.item() 
            accuracy=(outputs.argmax(1)==targets).sum() #预判正确
            total_accuracy+=accuracy
    print('total test accuracy {}'.format(total_accuracy/len(test_data))) #测试正确率
    print("total test loss {}".format(total_test_loss))
    writer.add_scalar("test_loss",total_test_loss,total_test_step)
    total_test_step+=1
    
    # 保存每次的模型
    torch.save(mynet,"./data/05 data/myNet/model_%d.pth"%i) #保存模型
    
writer.close()

3.模型验证套路

import torch
from PIL import Image

image_path="./data/imgs/dog.png"
image=Image.open(image_path)
image=image.convert("RGB")
print(image)

from torchvision import transforms
transform=transforms.Compose([transforms.Resize((32,32)),transforms.ToTensor()])
image=transform(image)
print(image.shape)


# 搭建神经网络
class myNet(nn.Module):
    def __init__(self):
        super(myNet,self).__init__()
        self.model=nn.Sequential(
            nn.Conv2d(3,32,5,1,2)
            ,nn.MaxPool2d(2)
            ,nn.Conv2d(32,32,5,1,2)
            ,nn.MaxPool2d(2)
            ,nn.Conv2d(32,64,5,1,2)
            ,nn.MaxPool2d(2)
            ,nn.Flatten() # 64x4x4
            ,nn.Linear(64*4*4,64)
            ,nn.Linear(64,10)
        )
    def forward(self,x):
        output=self.model(x)
        return output
    
# 使用第一种方法加载模型
model=torch.load("./data/05 data/myNet/model_1.pth")
image=torch.reshape(image,(-1,3,32,32))

# 在CPU上加载GPU训练的模型
model=torch.load("./data/05 data/myNet/model_1.pth",map_location=torch.device("cpu")) 

model.eval()
with torch.no_grad():
    output=model(image)
    
print(output)
print(output.argmax(1)) #打印训练出来的结果

参考资料

[1] b站课程链接

[2] 手敲代码链接