# coding = utf-8
import numpy as np
import torch
from torchvision import transforms
_task = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(
[0.5], [0.5]
)
])
from torchvision.datasets import MNIST
# 数据集加载
mnist = MNIST('./data', download=True, train=True, transform=_task)
# 训练集和验证集划分
from torch.utils.data import DataLoader
from torch.utils.data.sampler import SubsetRandomSampler
# create training and validation split
index_list = list(range(len(mnist)))
split_train = int(0.8 * len(mnist))
split_valid = int(0.9 * len(mnist))
train_idx, valid_idx, test_idx = index_list[:split_train], index_list[split_train:split_valid], index_list[split_valid:]
# create sampler objects using SubsetRandomSampler
train_sampler = SubsetRandomSampler(train_idx)
valid_sampler = SubsetRandomSampler(valid_idx)
test_sampler = SubsetRandomSampler(test_idx)
# create iterator objects for train and valid dataset
trainloader = DataLoader(mnist, batch_size=256, sampler=train_sampler)
validloader = DataLoader(mnist, batch_size=256, sampler=valid_sampler)
test_loader = DataLoader(mnist, batch_size=256, sampler=test_sampler)
# design for net
import torch.nn.functional as F
class NetModel(torch.nn.Module):
def __init__(self):
super(NetModel, self).__init__()
self.hidden = torch.nn.Linear(28 * 28, 300)
self.output = torch.nn.Linear(300, 10)
def forward(self, x):
x = x.view(-1, 28 * 28)
x = self.hidden(x)
x = F.relu(x)
x = self.output(x)
return x
if __name__ == "__main__":
net = NetModel()
from torch import optim
loss_function = torch.nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.01, weight_decay=1e-6, momentum=0.9, nesterov=True)
for epoch in range(1, 12):
train_loss, valid_loss = [], []
# net.train()
for data, target in trainloader:
optimizer.zero_grad()
# forward propagation
output = net(data)
loss = loss_function(output, target)
loss.backward()
optimizer.step()
train_loss.append(loss.item())
# net.eval()
for data, target in validloader:
output = net(data)
loss = loss_function(output, target)
valid_loss.append(loss.item())
print("Epoch:", epoch, "Training Loss:", np.mean(train_loss), "Valid Loss:", np.mean(valid_loss))
print("testing ... ")
total = 0
correct = 0
for i, test_data in enumerate(test_loader, 0):
data, label = test_data
output = net(data)
_, predict = torch.max(output.data, 1)
total += label.size(0)
correct += np.squeeze((predict == label).sum().numpy())
print("Accuracy:", (correct / total) * 100, "%")
pytorch2-基于minst数据集的简单全连接网络
最新推荐文章于 2021-07-10 16:27:30 发布
![](https://img-home.csdnimg.cn/images/20240711042549.png)