PyTorch 基础篇(4):前馈神经网络(Feedforward Neural Network)
参考代码
yunjey的 pytorch tutorial系列
1 2 3 4 5
| import torch import torch.nn as nn import torchvision import torchvision.transforms as transforms
|
1 2 3 4 5 6 7 8 9 10 11
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
input_size = 784 hidden_size = 500 num_classes = 10 num_epochs = 5 batch_size = 100 learning_rate = 0.001
|
cuda
MINIST 数据集加载
1 2 3 4 5
| train_dataset = torchvision.datasets.MNIST(root='../../../data/minist', train=True, transform=transforms.ToTensor(), download=True)
|
Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz
Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz
Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz
Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz
Processing...
Done!
1 2 3 4
| test_dataset = torchvision.datasets.MNIST(root='../../../data/minist', train=False, transform=transforms.ToTensor())
|
1 2 3 4 5 6 7 8 9 10
|
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)
|
自定义前馈神经网络
1 2 3 4 5 6 7 8 9 10 11 12 13
| class NeuralNet(nn.Module): def __init__(self, input_size, hidden_size, num_classes): super(NeuralNet, 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
|
1 2 3 4 5 6 7 8
|
model = NeuralNet(input_size, hidden_size, num_classes).to(device)
criterion = nn.CrossEntropyLoss() optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
|
训练模型
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20
| total_step = len(train_loader)
for epoch in range(num_epochs): for i, (images, labels) in enumerate(train_loader): images = images.reshape(-1, 28*28).to(device) labels = labels.to(device) outputs = model(images) loss = criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() if (i+1) % 100 == 0: print ('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}' .format(epoch+1, num_epochs, i+1, total_step, loss.item()))
|
Epoch [1/5], Step [100/600], Loss: 0.3552
Epoch [1/5], Step [200/600], Loss: 0.2119
Epoch [1/5], Step [300/600], Loss: 0.3251
Epoch [1/5], Step [400/600], Loss: 0.2134
Epoch [1/5], Step [500/600], Loss: 0.1376
Epoch [1/5], Step [600/600], Loss: 0.1637
Epoch [2/5], Step [100/600], Loss: 0.1140
Epoch [2/5], Step [200/600], Loss: 0.2439
Epoch [2/5], Step [300/600], Loss: 0.1156
Epoch [2/5], Step [400/600], Loss: 0.0217
Epoch [2/5], Step [500/600], Loss: 0.0973
Epoch [2/5], Step [600/600], Loss: 0.1627
Epoch [3/5], Step [100/600], Loss: 0.1536
Epoch [3/5], Step [200/600], Loss: 0.0502
Epoch [3/5], Step [300/600], Loss: 0.0367
Epoch [3/5], Step [400/600], Loss: 0.1291
Epoch [3/5], Step [500/600], Loss: 0.0502
Epoch [3/5], Step [600/600], Loss: 0.0670
Epoch [4/5], Step [100/600], Loss: 0.0598
Epoch [4/5], Step [200/600], Loss: 0.0823
Epoch [4/5], Step [300/600], Loss: 0.0466
Epoch [4/5], Step [400/600], Loss: 0.0350
Epoch [4/5], Step [500/600], Loss: 0.0754
Epoch [4/5], Step [600/600], Loss: 0.0601
Epoch [5/5], Step [100/600], Loss: 0.0274
Epoch [5/5], Step [200/600], Loss: 0.0469
Epoch [5/5], Step [300/600], Loss: 0.1103
Epoch [5/5], Step [400/600], Loss: 0.0505
Epoch [5/5], Step [500/600], Loss: 0.0093
Epoch [5/5], Step [600/600], Loss: 0.0513
测试并保存模型
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
|
with torch.no_grad(): correct = 0 total = 0 for images, labels in test_loader: images = images.reshape(-1, 28*28).to(device) labels = labels.to(device) outputs = model(images) _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item()
print('Accuracy of the network on the 10000 test images: {} %'.format(100 * correct / total))
|
Accuracy of the network on the 10000 test images: 97.47 %
1 2
| torch.save(model.state_dict(), 'model.ckpt')
|