CNN实现MNIST数据集分类时,发现CPU训练特别慢,20个epoch需要半个多小时,故修改为GPU训练,训练时长明显缩短。
修改如下:
-
需要把网络结构移到GPU上:
model = Net() if torch.cuda.is_available(): print("The model is using GPU") model = model.cuda()
-
训练数据集和测试数据集移到GPU上:
for i, data in enumerate(train_loader): # 获得数据和对应的标签 inputs, labels = data # 获得模型预测结果,(64,10) inputs = inputs.cuda() labels = labels.cuda() out = model(inputs) # 交叉熵代价函数out(batch,C),labels(batch) loss = entropy_loss(out, labels)
注意: 刚开始由于只修改了数据,未修改标签,导致报错:
RuntimeError: Expected object of device type cuda but got device type cpu for argument #2 ‘mat1’ in call to _th_addmm
完整代码如下——使用工具:jupyter notebook:
#%%
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
#%%
# 下载训练集
train_dataset = datasets.MNIST(root='./',
train=True,
transform=transforms.ToTensor(),
download=False)
# 下载测试集
test_dataset = datasets.MNIST(root='./',
train=False,
transform=transforms.ToTensor(),
download=False)
#%%
# 批次大小
batch_size = 64
# 装载训练集
train_loader = DataLoader(dataset=train_dataset,
batch_size=batch_size,
shuffle=True)
# 装载测试集
test_loader = DataLoader(dataset=test_dataset,
batch_size=batch_size,
shuffle=True)
#%%
for i, data in enumerate(train_loader):
# 获得数据和对应的标签
inputs, labels = data
print(inputs.shape)
print(labels.shape)
break
# 打印部分数据集图片
for i in range(2):
inputs,labels = train_dataset[i]
print(labels)
plt.imshow(inputs.reshape(28,28).numpy(),cmap='gray')
plt.show()
#%%
# 定义网络结构
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
# Conv2d(in_channel,out_channel(32个特征图),kernel_size,stride,padding(补几圈零))
# 如果卷积窗口是3*3,则padding为1
# 如果是5*5,则填充两圈 7*7填充3
# 如果要增加nn.Dropout(p=0.5) p为丢失的概率
self.conv1 = nn.Sequential(nn.Conv2d(1, 32, 5, 1, 2), nn.ReLU(), nn.MaxPool2d(2, 2))
self.conv2 = nn.Sequential(nn.Conv2d(32, 64, 5, 1, 2), nn.ReLU(), nn.MaxPool2d(2, 2))
# 64个特征图 特征图为7*7 因为28*28做了两次池化
self.fc1 = nn.Sequential(nn.Linear(64 * 7 * 7, 1000), nn.Dropout(p=0.4), nn.ReLU())
self.fc2 = nn.Sequential(nn.Linear(1000, 10), nn.Softmax(dim=1))
def forward(self, x):
# ([64, 1, 28, 28]) 卷积对格式有要求,必须为4维 batch_size channel 图片大小
x = self.conv1(x)
x = self.conv2(x)
# 全连接为2维数据 (64,1,7,7)——>
x = x.view(x.size()[0], -1)
x = self.fc1(x)
x = self.fc2(x)
return x
#%%
LR = 0.0003
# 定义模型 把神经网络移到GPU上
model = Net()
if torch.cuda.is_available():
print("The model is using GPU")
model = model.cuda()
# 定义代价函数
entropy_loss = nn.CrossEntropyLoss()
# 定义优化器 Adam比SGD好一点
optimizer = optim.Adam(model.parameters(), LR)
#%%
def train():
model.train()
for i, data in enumerate(train_loader):
# 获得数据和对应的标签
inputs, labels = data
# 获得模型预测结果,(64,10)
inputs = inputs.cuda()
labels = labels.cuda()
out = model(inputs)
# 交叉熵代价函数out(batch,C),labels(batch)
loss = entropy_loss(out, labels)
# 梯度清0
optimizer.zero_grad()
# 计算梯度
loss.backward()
# 修改权值
optimizer.step()
def test():
model.eval()
correct = 0
for i, data in enumerate(test_loader):
# 获得数据和对应的标签
inputs, labels = data
# 获得模型预测结果
inputs = inputs.cuda()
labels = labels.cuda()
out = model(inputs)
# 获得最大值,以及最大值所在的位置
_, predicted = torch.max(out, 1)
# 预测正确的数量
correct += (predicted == labels).sum()
print("Test acc: {0}".format(correct.item() / len(test_dataset)))
correct = 0
for i, data in enumerate(train_loader):
# 获得数据和对应的标签
inputs, labels = data
# 获得模型预测结果
inputs = inputs.cuda()
labels = labels.cuda()
out = model(inputs)
# 获得最大值,以及最大值所在的位置
_, predicted = torch.max(out, 1)
# 预测正确的数量
correct += (predicted == labels).sum()
print("Train acc: {0}".format(correct.item() / len(train_dataset)))
#%%
for epoch in range(0, 20):
print('epoch:',epoch)
train()
test()
#%%