经过一段时间的学习,终于初步入门了pytorch,经过这么长时间的学习,收获了很多,我们一定要记住,天道酬勤,学习一定要坚持不懈,终身学习,养成学习的兴趣!加油,我们都能在生活中找到自己的目标!今天是学习笔记入门的最后一篇,主要写了将神经网络模型送入gpu进行运算的代码,同时也写了验证的代码,大家一起加油!
import torch.optim
import torchvision
import time
#准备数据集
from tensorboardX import SummaryWriter
from torch import nn
from torch.utils.data import DataLoader
# from model import *
#准备训练设备
device = torch.device("cuda")
train_data = torchvision.datasets.CIFAR10(root="dataset",train=True,transform=torchvision.transforms.ToTensor(),
download=True)
test_data = torchvision.datasets.CIFAR10(root="dataset",train=False,transform=torchvision.transforms.ToTensor(),
download=True)
# length 长度
train_data_size = len(train_data)
test_data_size = len(test_data)
# 如果train_data_size = 10,训练数据集的长度为:10
print("训练数据集的长度为:{}".format(train_data_size))
print("测试数据集的长度为:{}".format(test_data_size))
# 利用 DataLoader 来加载数据集
train_dataloader = DataLoader(train_data,batch_size=64)
test_dataloader = DataLoader(test_data,batch_size=64)
#创建神经网络
class B(nn.Module):
def __init__(self):
super(B, 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(),
nn.Linear(64*4*4,64),
nn.Linear(64,10)
)
def forward(self,x):
x = self.model(x)
return x
b = B()
b = b.to(device)
#损失函数
loss_fn = nn.CrossEntropyLoss()
loss_fn = loss_fn.to(device)
#优化器
learning_rate = 1e-2
optimizer = torch.optim.SGD(b.parameters(),lr=learning_rate)
#设置训练网络的一些参数
#记录模型参数
i=0
#记录测试的次数
total_test_step =0
#记录训练的次数
total_train_step = 0
#训练的轮数
epoch = 20
#添加tensorboard
writer = SummaryWriter("logs")
start_time = time.time()
for i in range(epoch):
print("-------第{}轮训练开始-------".format(i+1))
#训练步骤开始
b.train()
for data in train_dataloader:
imgs,targets = data
imgs = imgs.to(device)
targets = targets.to(device)
output = b(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("训练次数:{},Loss:{}".format(total_train_step,loss.item()))
writer.add_scalar("train_loss",loss.item(),total_train_step)
#测试步骤开始
b.eval()
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 = b(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("整体测试集上的Loss:{}".format(total_test_loss))
print("整体测试集上的正确率:{}".format(total_accuracy/test_data_size))
writer.add_scalar("test_loss",total_test_loss,total_test_step)
writer.add_scalar("test_accuracy",total_accuracy/test_data_size,total_test_step)
total_test_step = total_test_step+1
torch.save(b,"tudui_{}.pth".format(i))
print("模型保存成功")
i=i+1
下面是验证的代码:
import torch
import torchvision
from PIL import Image
from torch import nn
image = Image.open("imgs/bird2.jpg")
transform = torchvision.transforms.Compose([
torchvision.transforms.Resize((32,32)),
torchvision.transforms.ToTensor()
])
image = transform(image)
class B(nn.Module):
def __init__(self):
super(B, 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(),
nn.Linear(64*4*4,64),
nn.Linear(64,10)
)
def forward(self,x):
x = self.model(x)
return x
model = torch.load("tudui_19.pth",map_location='cpu')
print(model)
image = torch.reshape(image,(1,3,32,32))
model.eval()
with torch.no_grad():
output = model(image)
print(output)
print(output.argmax(1))