【小土堆】 30-31/33 利用GPU训练模型

有两种训练方法

第一种  调用cuda（）方法

代码实现：

import torch
import torchvision
from torch import nn
from torch.utils.tensorboard import SummaryWriter
#引入网络模型文件
# from model import  *
from torch.utils.data import DataLoader
import time

#准备训练数据集和测试数据集
train_data = torchvision.datasets.CIFAR10(root='./data', train=True, download=True,
                                          transform=torchvision.transforms.ToTensor())
test_data = torchvision.datasets.CIFAR10(root='./data', train=False, download=True,
                                         transform=torchvision.transforms.ToTensor())

#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_loader = DataLoader(train_data, batch_size=64, shuffle=True)
test_loader = DataLoader(test_data, batch_size=64, shuffle=True)


#网络模型  （单独的model.py 文件）
class Tudui(nn.Module):
    def __init__(self):
        super(Tudui, self).__init__()
        self.model1 = nn.Sequential(
        nn.Conv2d(in_channels=3, out_channels=32, kernel_size=5, padding=2),
        nn.MaxPool2d(kernel_size=2),
        nn.Conv2d(in_channels=32, out_channels=32, kernel_size=5, padding=2),
        nn.MaxPool2d(kernel_size=2),
        nn.Conv2d(in_channels=32, out_channels=64, kernel_size=5, padding=2),
        nn.MaxPool2d(kernel_size=2),
        # 展现数据
        nn.Flatten(),
        nn.Linear(in_features=1024, out_features=64),
        nn.Linear(in_features=64, out_features=10)
        )

    def forward(self, x):
        x = self.model1(x)
        return x
tudui = Tudui()
tudui = tudui.cuda()
#损失函数
loss_fu = nn.CrossEntropyLoss()
loss_fu = loss_fu.cuda()
#优化器优化模型
#1e-2 = (10)^(-2) = 1/100 =0.01
learing_rate = 1e-2
optimizer = torch.optim.SGD(tudui.parameters(), lr=learing_rate)

#设置训练网络的一些参数
#记录训练的次数
total_train_step = 0
#记录测试的次数
total_test_step = 0
#训练轮数
epoch = 10
#添加tensorboard
writer = SummaryWriter('./logs_train')
for i in range(epoch):
    print("第{}轮训练开始".format(i+1))
    start_time = time.time()
    # 训练开始
    for data in train_loader:
        imgs, targets = data
        imgs, targets = imgs.cuda(), targets.cuda()
        ouputs = tudui(imgs)
        loss = loss_fu(ouputs, targets)
        #优化器 优化模型
        #梯度清零
        optimizer.zero_grad()
        loss.backward()
        optimizer.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)

    #测试步骤开始
    total_test_loss = 0.0
    #正确率
    total_accuracy = 0.0
    with torch.no_grad():
        for data in test_loader:
            imgs,targets = data
            imgs, targets = imgs.cuda(), targets.cuda()
            outputs = tudui(imgs)
            loss = loss_fu(outputs, targets)
            total_test_loss += loss.item()
            accuracy = (outputs.argmax(1) == targets).sum()
            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 += 1
    #保存模型
    torch.save(tudui, 'tudui_{}.pth'.format(i))
    print("模型已保存")
    writer.close()

第二种  更常用的一种方式，调用.to

先定义训练设备

#定义训练设备
# device = torch.device("cpu")
device = torch.device("cuda:0")

然后调用.to（device）

整体代码：

import torch
import torchvision
from torch import nn
from torch.utils.tensorboard import SummaryWriter
#引入网络模型文件
# from model import  *
from torch.utils.data import DataLoader
import time

#定义训练设备
# device = torch.device("cpu")
device = torch.device("cuda:0")
#准备训练数据集和测试数据集
train_data = torchvision.datasets.CIFAR10(root='./data', train=True, download=True,
                                          transform=torchvision.transforms.ToTensor())
test_data = torchvision.datasets.CIFAR10(root='./data', train=False, download=True,
                                         transform=torchvision.transforms.ToTensor())

#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_loader = DataLoader(train_data, batch_size=64, shuffle=True)
test_loader = DataLoader(test_data, batch_size=64, shuffle=True)


#网络模型  （单独的model.py 文件）
class Tudui(nn.Module):
    def __init__(self):
        super(Tudui, self).__init__()
        self.model1 = nn.Sequential(
        nn.Conv2d(in_channels=3, out_channels=32, kernel_size=5, padding=2),
        nn.MaxPool2d(kernel_size=2),
        nn.Conv2d(in_channels=32, out_channels=32, kernel_size=5, padding=2),
        nn.MaxPool2d(kernel_size=2),
        nn.Conv2d(in_channels=32, out_channels=64, kernel_size=5, padding=2),
        nn.MaxPool2d(kernel_size=2),
        # 展现数据
        nn.Flatten(),
        nn.Linear(in_features=1024, out_features=64),
        nn.Linear(in_features=64, out_features=10)
        )

    def forward(self, x):
        x = self.model1(x)
        return x
tudui = Tudui()
tudui = tudui.to(device)
#损失函数
loss_fu = nn.CrossEntropyLoss()
loss_fu = loss_fu.to(device)
#优化器优化模型
#1e-2 = (10)^(-2) = 1/100 =0.01
learing_rate = 1e-2
optimizer = torch.optim.SGD(tudui.parameters(), lr=learing_rate)

#设置训练网络的一些参数
#记录训练的次数
total_train_step = 0
#记录测试的次数
total_test_step = 0
#训练轮数
epoch = 10
#添加tensorboard
writer = SummaryWriter('./logs_train')
for i in range(epoch):
    print("第{}轮训练开始".format(i+1))
    start_time = time.time()
    # 训练开始
    for data in train_loader:
        imgs, targets = data
        imgs, targets = imgs.to(device), targets.to(device)
        ouputs = tudui(imgs)
        loss = loss_fu(ouputs, targets)
        #优化器 优化模型
        #梯度清零
        optimizer.zero_grad()
        loss.backward()
        optimizer.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)

    #测试步骤开始
    total_test_loss = 0.0
    #正确率
    total_accuracy = 0.0
    with torch.no_grad():
        for data in test_loader:
            imgs,targets = data
            imgs, targets = imgs.to(device), targets.to(device)
            outputs = tudui(imgs)
            loss = loss_fu(outputs, targets)
            total_test_loss += loss.item()
            accuracy = (outputs.argmax(1) == targets).sum()
            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 += 1
    #保存模型
    torch.save(tudui, 'tudui_{}.pth'.format(i))
    print("模型已保存")
    writer.close()