使用GPU训练小型网络，以训练CIFAR10数据集为例

          使用GPU仅可以通过修改设备，将可使用的GPU加进来，载加载数据集时使用GPU来加载，计算网络时，在构建计算图时，将以前用cpu的形式改为GPU即可

第一步：先测试是否有GPU或GPU是否可用，不可用可能是你的配置有问题，例如CUDA，pytorch等版本与你的显卡不匹配的原因，测试：

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

正常情况下可用会输出：cuda 0,不可用则输出cpu

如果配置正常可是尝试cuda 1是否可行

第二步：将相应部分的设备改为GPU即可

①网络  

net=net.to(device)

②损失函数以我的交叉熵为例

Loss=nn.CrossEntropyLoss()#交叉熵
Loss=Loss.to(device)

③数据加载

将图片和标签通过GPU加载

imgs=imgs.to(device)
targets=targets.to(device)

以上如果有错误请求指教改正，入门学生谢过！

 以下为测试例子，训练数据集为CIFAR10数据集的分类问题，使用的网络为LeNet

import torch.optim
import torchvision
from torch import nn
from torch.utils.data import DataLoader


# 准备数据集合
# from torch.utils.tensorboard import SummaryWriter

train_data=torchvision.datasets.CIFAR10(
    root="../data",train=True,
    transform=torchvision.transforms.ToTensor(),
    download=False
)
test_data=train_test=torchvision.datasets.CIFAR10(
    root="../data",train=False,
    transform=torchvision.transforms.ToTensor(),
    download=False
)
train_size=len(train_data)
test_size=len(test_data)
# print(train_size)
# print(test_size)
train_loader=DataLoader(dataset=train_data,batch_size=64)
test_loader=DataLoader(dataset=test_data,batch_size=64)
#
# 2.搭建神经网络
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.module=nn.Sequential(
            nn.Conv2d(3,32,5,padding=2),
            nn.MaxPool2d(2),
            nn.Conv2d(32,32,5,padding=2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 5, padding=2),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(1024,64),
            nn.Linear(64,10)
        )
    def forward(self,input):
        output=self.module(input)
        return output
net=Net()

# gpu

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

# net=Net.cuda()
# 3.构造优化器
Loss=nn.CrossEntropyLoss()#交叉熵
# Gpu
# Loss=Loss.cuda()

Loss=Loss.to(device)
optimizaer=torch.optim.SGD(net.parameters(),lr=0.01)

###############引入tensorboard
# writer=SummaryWriter("./train_log")

# 4.训练

test_count=0
total_count=0


for e in range(5):
    print(f"_______第{e+1}轮训练开始_______")
    for data in train_loader:
        imgs,targets=data

        # Gpu
        imgs=imgs.to(device)
        targets=targets.to(device)

        # 前馈
        output=net(imgs)
        loss=Loss(output,targets)
        # 反馈
        optimizaer.zero_grad()
        loss.backward()
        optimizaer.step()
        test_count+=1
        if test_count%100==0:
            # writer.add_scalar("train_loss",loss.item(),test_count)
            print(f"训练次数:{test_count+1}，Loss:{loss.item()}")
    # 测试步骤开始
    sum=0
    with torch.no_grad():
        for data in test_loader:
            imgs, targets = data
            # Gpu
            imgs = imgs.to(device)
            targets = targets.to(device)
            output = net(imgs)
            loss = Loss(output,targets)
            sum=sum+loss
    total_count+=1
    print(f"第{e+1}轮训练整体的Loss为{sum}")    # writer.add_scalar("test_loss",sum,total_count)
    torch.save(net,f"net_{e+1}.pth")
    print(f"第{e+1}轮模型以保存")



# writer.close()