pytorch测试demo
pytorch一个简单的测试代码。
手写数字识别
😁😁😁😁
# coding: utf-8
import torch
import torch.nn as nn
import torchvision.datasets as dsets
import torchvision.transforms as transforms
from torch.autograd import Variable
import torch.utils.data as Data
import matplotlib.pyplot as plt
input_size = 784 # 输入大小
hidden_size = 500 # 第一层的神经元个数
num_classes = 10 # 类别个数
num_epochs = 5 # 迭代次数
batch_size = 100 # 每一批的数目
learning_rate = 0.001 # 学习率
train_dataset = dsets.MNIST(root='./data', train=True,
transform=transforms.ToTensor(), download=True) # 加载数据集
test_dataset = dsets.MNIST(root='./data', train=False,
transform=transforms.ToTensor(), download=True) # 加载测试集
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
batch_size=batch_size, shuffle=True) # shuffle为True:在每个epoch重新排列数据
test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
batch_size=batch_size, shuffle=False)
# In[4]:
class Net(nn.Module):
def __init__(self, input_size, hidden_size, num_classes):
super(Net, self).__init__()
self.fc1 = nn.Linear(input_size, hidden_size)
self.relu = nn.ReLU()
self.fc2 = nn.Linear(hidden_size, num_classes)
def forward(self, x):
out = self.fc1(x)
out = self.relu(out)
out = self.fc2(out)
return out
# In[5]:
net = Net(input_size, hidden_size, num_classes) # 创建模型
criterion = nn.CrossEntropyLoss() # 创建损失函数(用的交叉熵)
optimizer = torch.optim.Adam(net.parameters(), lr=learning_rate) # 梯度下降
# In[6]:
for epoch in range(num_epochs):
for i, (images, labels) in enumerate(train_loader):
images = Variable(images.view(-1, 28 * 28))
label = Variable(labels)
optimizer.zero_grad() # 将权重初始化为0
outputs = net(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step() # 更新权重
if (i + 1) % 100 == 0:
print('Epoch [%d/%d], Step [%d/%d],Loss:%.4f' % (epoch + 1, num_epochs, i + 1,
len(train_dataset) // batch_size, loss.item()))
# In[7]:
# 预测数据
correct = 0
total = 0
for images, labels in test_loader:
images = Variable(images.view(-1, 28 * 28))
outputs = net(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum()
print('Accuracy of the network on the 10000 test images:%d %%' % (100 * correct // total))
# In[8]:
plt.imshow(test_dataset.data[-25].numpy(), cmap='gray')
plt.title('%i ' % (predicted[-25]))
plt.show()
# In[9]:
# 储存模型
torch.save(net, 'model1.pkl')
# In[10]:
# 加载模型
model = torch.load('model1.pkl')
# In[11]:
# 预测数据
correct = 0
total = 0
for images, labels in test_loader:
images = Variable(images.view(-1, 28 * 28))
outputs = model(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum()
print('Accuracy of the network on the 10000 test images:%d %%' % (100 * torch.true_divide(correct, total)))
# In[12]:
torch.save(net.state_dict(), 'model2.pkl')
net.load_state_dict(torch.load('model2.pkl'))
评论
隐私政策
你无需删除空行,直接评论以获取最佳展示效果