是参考B站一位up主的,感觉讲的很好,讲解的每节课程在csdn上做了一个目录: https://blog.csdn.net/qq_37541097/article/details/103482003
这一节的up主的参考资料如下:
这个例子是在PyTorch官网上讲解的:
中文翻译版可以参考:
01 运行demo
1.1 文件结构
-
model.py ——定义LeNet网络模型 -
train.py ——加载数据集并训练,训练集计算loss,测试集计算accuracy,保存训练好的网络参数 -
predict.py ——得到训练好的网络参数后,用自己找的图像进行分类测试
1.2 下载数据集
import torchvision
import torchvision.transforms as transforms
transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
train_set = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
运行结果:
1.3 代码
model.py
import torch.nn as nn
import torch.nn.functional as F
class LeNet(nn.Module):
def __init__(self):
super(LeNet, self).__init__()
self.conv1 = nn.Conv2d(3, 16, 5)
self.pool1 = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(16, 32, 5)
self.pool2 = nn.MaxPool2d(2, 2)
self.fc1 = nn.Linear(32*5*5, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
def forward(self, x):
x = F.relu(self.conv1(x))
x = self.pool1(x)
x = F.relu(self.conv2(x))
x = self.pool2(x)
x = x.view(-1, 32*5*5)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
train.py
import torch
import torchvision
import torch.nn as nn
from model import LeNet
import torch.optim as optim
import torchvision.transforms as transforms
def main():
transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
train_set = torchvision.datasets.CIFAR10(root='./data', train=True,
download=False, transform=transform)
train_loader = torch.utils.data.DataLoader(train_set, batch_size=36,
shuffle=True, num_workers=0)
val_set = torchvision.datasets.CIFAR10(root='./data', train=False,
download=False, transform=transform)
val_loader = torch.utils.data.DataLoader(val_set, batch_size=5000,
shuffle=False, num_workers=0)
val_data_iter = iter(val_loader)
val_image, val_label = val_data_iter.next()
net = LeNet()
loss_function = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr=0.001)
for epoch in range(5):
running_loss = 0.0
for step, data in enumerate(train_loader, start=0):
inputs, labels = data
optimizer.zero_grad()
outputs = net(inputs)
loss = loss_function(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
if step % 500 == 499:
with torch.no_grad():
outputs = net(val_image)
predict_y = torch.max(outputs, dim=1)[1]
accuracy = torch.eq(predict_y, val_label).sum().item() / val_label.size(0)
print('[%d, %5d] train_loss: %.3f test_accuracy: %.3f' %
(epoch + 1, step + 1, running_loss / 500, accuracy))
running_loss = 0.0
print('Finished Training')
save_path = './Lenet.pth'
torch.save(net.state_dict(), save_path)
if __name__ == '__main__':
main()
predict.py
import torch
import torchvision.transforms as transforms
from PIL import Image
from model import LeNet
def main():
transform = transforms.Compose(
[transforms.Resize((32, 32)),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
classes = ('plane', 'car', 'bird', 'cat',
'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
net = LeNet()
net.load_state_dict(torch.load('Lenet.pth'))
im = Image.open('1.jpg')
im = transform(im)
im = torch.unsqueeze(im, dim=0)
with torch.no_grad():
outputs = net(im)
predict = torch.max(outputs, dim=1)[1].data.numpy()
print(classes[int(predict)])
if __name__ == '__main__':
main()
1.4 预测
- 先运行train.py生成网络。运行结果如下:
好像PyTorch1.90有点问题,所以会报下面的警告,但不影响正常使用。
D:\application\anaconda\envs\mytorch\python.exe D:/study/mytorch/train.py
D:\application\anaconda\envs\mytorch\lib\site-packages\torch\nn\functional.py:718: UserWarning: Named tensors and all their associated APIs are an experimental feature and subject to change. Please do not use them for anything important until they are released as stable. (Triggered internally at ..\c10/core/TensorImpl.h:1156.)
return torch.max_pool2d(input, kernel_size, stride, padding, dilation, ceil_mode)
02 代码说明
关于代码的说明基本都已经在上面的代码中注释了。
其他还可以参考下面的两个:
想查看数据集中的图片可以使用下述代码: 需要注意的是, pytorch的数据张量形式为:[batch, channel, height, width] 而图片的数据形式为:[height,width,channel] 所以图片转pytorch中张量,就需要换数据形式,并且在最前面扩增一维; pytorch中张量转图片,也需要话数据顺序,比如plt.imshow(np.tr anspose(npimg, (1, 2, 0)))
03 使用gpu运行train.py
import torch
import torchvision
import torch.nn as nn
from model import LeNet
import torch.optim as optim
import torchvision.transforms as transforms
def main():
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)
transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
train_set = torchvision.datasets.CIFAR10(root='./data', train=True,
download=False, transform=transform)
train_loader = torch.utils.data.DataLoader(train_set, batch_size=36,
shuffle=True, num_workers=0)
val_set = torchvision.datasets.CIFAR10(root='./data', train=False,
download=False, transform=transform)
val_loader = torch.utils.data.DataLoader(val_set, batch_size=5000,
shuffle=False, num_workers=0)
val_data_iter = iter(val_loader)
val_image, val_label = val_data_iter.next()
net = LeNet()
net.to(device)
loss_function = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr=0.001)
for epoch in range(5):
running_loss = 0.0
for step, data in enumerate(train_loader, start=0):
inputs, labels = data
optimizer.zero_grad()
outputs = net(inputs.to(device))
loss = loss_function(outputs, labels.to(device))
loss.backward()
optimizer.step()
running_loss += loss.item()
if step % 500 == 499:
with torch.no_grad():
outputs = net(val_image.to(device))
predict_y = torch.max(outputs, dim=1)[1]
accuracy = (predict_y == val_label.to(device)).sum().item() / val_label.size(0)
print('[%d, %5d] train_loss: %.3f test_accuracy: %.3f' %
(epoch + 1, step + 1, running_loss / 500, accuracy))
running_loss = 0.0
print('Finished Training')
save_path = './Lenet.pth'
torch.save(net.state_dict(), save_path)
if __name__ == '__main__':
main()
说明
- 定义device为gpu
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
或直接指定
device = torch.device("cuda")
- 需要用
to() 函数来将Tensor在CPU和GPU之间相互移动,分配到指定的device中计算
① 定义网络后,将网络分配到gpu
net = LeNet()
net.to(device)
② 将训练过程中的inputs和labels分配到gpu中
outputs = net(inputs.to(device))
loss = loss_function(outputs, labels.to(device))
③ 将val_image分配到指定的gpu中
outputs = net(val_image.to(device))
accuracy = (predict_y == val_label.to(device)).sum().item() / val_label.size(0)
综上,定义device+把张量放入device,一共有6处需要修改
|