# feed forward nn
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.datasets as dsets
import torchvision.transforms as transforms
from torch.autograd import Variable
import argparse
parser = argparse.ArgumentParser(description='feedforward_mnist')
# learning & saving parameterss
parser.add_argument('-train', default = False, help = 'train the model')
parser.add_argument('-test', default = True, help = 'test the model')
parser.add_argument('-learning_rate', type = float, default = 0.005, help = 'initial learning rate [default = 0.001')
parser.add_argument('-num_epochs', type = int, default = 5, help = 'number of epochs of training [default = 10')
parser.add_argument('-batch_size', type = int, default = 100, help = 'batch size for training')
parser.add_argument('-input_size', type = int, default = 784, help = 'input size')
parser.add_argument('-hidden_size', type = int, default = 500, help = 'hidden size')
parser.add_argument('-output_size',type = int, default = 1, help = 'output size')
parser.add_argument('-num_classes', type = int, default = 10, help = 'hidden layer number')
parser.add_argument('-cuda',default = False, help = 'enable gpu')
args = parser.parse_args()
# Hyper Parameters
input_size = args.input_size
hidden_size = args.hidden_size
num_classes = args.num_classes
num_epochs = args.num_epochs
batch_size = args.batch_size
learning_rate = args.learning_rate
class Net(nn.Module):
def __init__(self, input_size, hidden_size, num_classes):
super(Net, self).__init__()
self.layer1 = nn.Linear(input_size, hidden_size)
# self.relu = nn.ReLU()
self.layer2 = nn.Linear(hidden_size, num_classes)
def forward(self, x):
x = F.leaky_relu(self.layer1(x))
x = self.layer2(x)
return x
net = Net(input_size, hidden_size, num_classes)
if args.train == True:
train_dataset = dsets.MNIST(root='./data',
train=args.train,
transform=transforms.ToTensor(),
download=True)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
batch_size=batch_size,
shuffle=True)
# Loss and Optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(net.parameters(), lr=learning_rate)
# Train the Model
for epoch in range(num_epochs):
for i, (images, labels) in enumerate(train_loader):
# Convert torch tensor to Variable
images = Variable(images.view(-1, 28*28))# image.size() = ([100, 1, 28, 28])
labels = Variable(labels)
# Forward + Backward + Optimize
optimizer.zero_grad() # zero the gradient buffer
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.data[0]))
torch.save(net.state_dict(), 'model.pkl')
# save only parameters
# MNIST Dataset
if args.test == True:
test_dataset = dsets.MNIST(root='./data',
train=args.train,
transform=transforms.ToTensor())
test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
batch_size=batch_size,
shuffle=False)
##############################
net.load_state_dict(torch.load('model.pkl'))
# load parameters to net model
##############################
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))
NN
©著作权归作者所有,转载或内容合作请联系作者
- 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
- 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
- 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
推荐阅读更多精彩内容
- Wilkinson GN, Eckert SR, Hancock TW, Mayo O. Nearest Neig...
- tf.nn.dropout(x, keep_prob, noise_shape=None, seed=None, ...
- 写在前面的话:哈喽,大家早安、午安、晚安喽,欢迎大家指点,也希望我的内容可以温暖、帮助同在学习路上的人们~ 正文开...