GPU加速 卷积网络MNIST模型训练

学习记录

import torch
from  model_mnist.Net import  ConvNet
from data.data_ import transform
from  torch.utils.data import  DataLoader
from torchvision.datasets import mnist
from torchvision import transforms
import time

MOMENTUM=0.5
LR = 0.01

transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize([0.1307],[0.3081])
     # 对于每个通道，从该通道的每个像素值中减去均值`0.1307`。
     #然后，将结果除以该通道的标准差`0.3081`。
     ]
)

if torch.cuda.is_available():
    print(f"Using GPU: {torch.cuda.get_device_name(0)}")  # 0是设备索引，通常是主设备
else:
    print("Using CPU")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


train_set = mnist.MNIST('./data', train=True, transform=transform, download=False)

train_loader = DataLoader(dataset = train_set,shuffle=True,batch_size=128,drop_last=True)

test_set = mnist.MNIST('./data', train=False, transform=transform, download=False)
test_loader = DataLoader(dataset = test_set,shuffle=True,batch_size=64,drop_last=True)


#实例化模型
model = ConvNet()
model.to(device)
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(),lr = LR, momentum=MOMENTUM)

def train(epoch):
    '''训练'''

    for index, data in enumerate(train_loader, 0):
        inputs, target = data# 取数据

        inputs = inputs.to(device)  # 将输入数据放到GPU上

        #print(inputs.device)
        target = target.to(device)  # 将标签放到GPU上
       #print(target.device)
        optimizer.zero_grad()  # 梯度清零
        outputs = model(inputs)  # 训练
        loss = criterion(outputs, target)  # 算损失
        loss.backward()  # 反向传播
        optimizer.step()  # 优化梯度

        # 累计loss ，这里要用item（）取数据要不回构建计算图
        if epoch% 20== 19:
            torch.save(model.state_dict(), "model_mnist/model_GPU.pth")  # 保存模型
            torch.save(optimizer.state_dict(), "model_mnist/optimize_GPUr.pth")
            print("训练次数为：{}，损失值为：{}".format(index, loss.item()))

if __name__ == '__main__':
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    start = time.time()
    for epoch in range(20):

        print({"————————第{}轮测试开始——————".format(epoch+ 1)})
        train(epoch)
    end = time.time()
    print('Time cost:', end - start, 's')

batch_size =64  训练20轮 CPU与GPU对比

bacth_size=128 训练20轮 CPU与GPU对比

GPU缩短了一半时间