首先需要准备数据集,本次使用的为公开数据集CIFAR10,其次设置好训练集和数据集
#导包
import torch
import torchvision
from torch.utils.tensorboard import SummaryWriter
from model import *
#准备数据集
from torch import nn
from torch.utils.data import DataLoader
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)
使用len函数,返回数据集的图片数量
#长度数据集
train_data_size = len(train_data)
test_data_size = len(test_data)
在print中使用format函数,进行数据格式化,并输出数据集的长度
#如果train_data_size=10,就会输出文本,本质上也是将数据格式化的手段
print("训练数据集的长度为:{}".format(train_data_size))
print("测试数据集的长度为:{}".format(test_data_size))
使用DataLoader函数进行数据集加载和读取功能。此函数块可填写,数据集名称和batch——size大小
#利用 DataLoader 加载数据集
train_dataloader = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)
创建网络模型,以tudui名为例,此网络模型,通过另一个python文件调用而来,调用命令:from xxx import*
tudui = Tudui()
设置损失函数
loss_fn = nn.CrossEntropyLoss()
设置优化器,使用SGD优化器,以下为两种lr损失函数表达方式
# 1e-2=1 x (10)^(-2) = 1/100=0.01
#第一种
# learning_rate = 0.01
#第二种
learning_rate = 1e-2
optimizer = torch.optim.SGD(tudui.parameters(), lr=learning_rate)
设置训练网络的参数
#记录训练的次数 从0开始
total_train_step = 0
#记录测试的次数
total_test_step = 0
#记录训练的轮数
epoch = 10
添加tensorboard命令,可以使用从功能直观显示损失函数的曲线图,调用方式:运行过后,在terminal控制台输入:tensorboard --logdir=xxx文件
#添加tensorboard
writer = SummaryWriter("../logs_train")
训练开始,使用for循环语句进行循环训练
for i in range(epoch):
print("-----第{}轮开始-----".format(i+1))
#训练步骤开始
# tudui.train() 此模式当模型有Dropout层时,会进行调用,可选择使用
for data in train_dataloader:
imgs ,targets = data
outputs = tudui(imgs)
loss = loss_fn(outputs, targets)
#优化器优化模型-梯度清零
optimizer.zero_grad()
#反向传播
loss.backward()
optimizer.step()
#结束一轮,继续+1训练
total_train_step = total_train_step +1
#将训练次数每100次显示一次
if total_train_step % 100 == 0:
print("训练次数{}, loss{}".format(total_train_step, loss.item( )))
writer.add_scalar("train_loss", loss.item(), total_train_step )
测试步骤开始
#测试步骤开始
# tudui.eval() 此模式当模型有Dropout层时,会进行调用,选择使用
total_test_loss = 0
total_accuracy = 0
with torch.no_grad():
for data in test_dataloader:
imgs, targets = data
outputs = tudui(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(tudui, "tudui_{}.pth".format(i))
print("模型已保存")
writer.close()
使用GPU进行模型训练
1.使用cuda语句:只可以在模型处、损失函数处、数据处:imgs/targets。这三个位置添加
#模型处
tudui = tudui.cuda()
#损失函数处
loss_fn = loss_fn.cuda()
#数据处
imgs = imgs.cuda()
targets = targets.cuda()
2.添加device和验证语句,选择是否使用gpu进行训练
#两种首部
#第一种
device = torch.device("cuda:0")/device("cuda")
#第二种
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)
#模型处
tudui = tudui.to(device)
#损失函数处
loss_fn = loss_fn.to(device)
#数据处
imgs = imgs.to(device)
targets = targets.to(device)