-
Notifications
You must be signed in to change notification settings - Fork 0
/
NetworkPytorch.py
88 lines (76 loc) · 2.95 KB
/
NetworkPytorch.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
import torch
import torch.nn as nn
import torchvision
from torch.autograd import Variable
import torch.utils.data as Data
class Net(nn.Module):
"""
构建网络模型
input_size = 28 * 28
hidden_size = 500:ReLU
num_classes = 10:Softmax(dim=1)
loss CrossEntropyLoss
"""
def __init__(self):
super().__init__()
self.fc1 = nn.Linear(28 * 28, 500)
self.relu = nn.ReLU()
self.fc2 = nn.Linear(500, 10)
def forward(self, x):
out = self.fc1(x)
out = self.relu(out)
out = self.fc2(out)
out = self.softmax(out)
return out
def softmax(self, x):
m = nn.Softmax(dim=1)
return m(x)
class LinearRegression(nn.Module):
def __init__(self, input_size, output_size):
super(LinearRegression, self).__init__()
self.linear = nn.Linear(input_size, output_size)
def forward(self, x):
out = self.linear(x)
return out
if __name__ == "__main__":
batch_size = 64
learning_rate = 0.001
# 加载数据
train_dataset = torchvision.datasets.MNIST(
root=r'D:\深度学习数据', train=True, download=True, transform=torchvision.transforms.ToTensor())
test_dataset = torchvision.datasets.MNIST(
root=r'D:\深度学习数据', train=False, download=True, transform=torchvision.transforms.ToTensor())
# 导入数据
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(
test_dataset, batch_size=batch_size, shuffle=True)
net = Net()
criterion = nn.CrossEntropyLoss() # 交叉熵损失函数
optimizer = torch.optim.Adam(net.parameters(), lr=learning_rate)
num_epochs = 5
for epoch in range(num_epochs):
for i, (images, labels) in enumerate(train_loader):
images = Variable(images.view(-1, 28 * 28)) # 声明变量用于学习参数。view相当于reshape
labels = Variable(labels)
optimizer.zero_grad() # 梯度清零
outputs = net(images)
loss = criterion(outputs, labels)
loss.backward() # 前向传播
optimizer.step() # 反向传播
# 打印指标
if (i + 1) % 100 == 0:
print('Epoch: [{}/{} ] Step [{}/{}], Loss: {:.4f}'.format(
epoch + 1, num_epochs, i + 1, len(train_dataset) // batch_size,
loss.item()
))
# 测试集验证准确率
correct = 0
total = 0
for images, labels in train_loader:
images = Variable(images.view(-1, 28 * 28)) # 声明变量用于学习参数。view相当于reshape
outputs = net(images)
_, predicted = torch.max(outputs.data, dim=1) # value, index
total += labels.size(0) # 多少数据
correct += (predicted == labels).sum()
print('Accuracy of the network on the 10000 test images:{:.4f}%'.format(100*correct/total))