|
import torch |
|
import torchvision |
|
from torch import optim |
|
import torch.nn.functional as F |
|
from torch.utils.data import DataLoader |
|
from Pytorch_MNIST图片识别.model import Net |
|
import matplotlib.pyplot as plt |
|
import os |
|
|
|
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" |
|
|
|
|
|
n_epochs = 3 |
|
|
|
|
|
batch_size_train = 64 |
|
batch_size_test = 1000 |
|
|
|
|
|
learning_rate = 0.01 |
|
momentum = 0.5 |
|
|
|
log_interval = 10 |
|
random_seed = 1 |
|
torch.manual_seed(random_seed) |
|
|
|
|
|
train_loader = torch.utils.data.DataLoader( |
|
torchvision.datasets.MNIST('./data/', train=True, download=True, |
|
transform=torchvision.transforms.Compose([ |
|
|
|
torchvision.transforms.ToTensor(), |
|
|
|
torchvision.transforms.Normalize( |
|
(0.1307,), (0.3081,)) |
|
])), |
|
batch_size=batch_size_train, shuffle=True) |
|
test_loader = torch.utils.data.DataLoader( |
|
torchvision.datasets.MNIST('./data/', train=False, download=True, |
|
transform=torchvision.transforms.Compose([ |
|
torchvision.transforms.ToTensor(), |
|
|
|
torchvision.transforms.Normalize( |
|
(0.1307,), (0.3081,)) |
|
])), |
|
batch_size=batch_size_test, shuffle=False) |
|
|
|
|
|
network = Net() |
|
optimizer = optim.SGD(network.parameters(), lr=learning_rate, momentum=momentum) |
|
|
|
train_losses = [] |
|
train_counter = [] |
|
test_losses = [] |
|
test_counter = [i * len(train_loader.dataset) for i in range(n_epochs + 1)] |
|
|
|
|
|
|
|
model_path = './model1/model.pth' |
|
optimizer_path = './model1/optimizer.pth' |
|
|
|
|
|
def train(epoch): |
|
network.train() |
|
for batch_idx, (data, target) in enumerate(train_loader): |
|
|
|
optimizer.zero_grad() |
|
output = network(data) |
|
loss = F.nll_loss(output, target) |
|
loss.backward() |
|
optimizer.step() |
|
if batch_idx % log_interval == 0: |
|
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( |
|
epoch, batch_idx * len(data), len(train_loader.dataset), 100. * batch_idx / len(train_loader), loss.item())) |
|
train_losses.append(loss.item()) |
|
train_counter.append( |
|
(batch_idx * 64) + ((epoch - 1) * len(train_loader.dataset))) |
|
if epoch == (n_epochs - 1): |
|
|
|
torch.save(network.state_dict(), model_path) |
|
torch.save(optimizer.state_dict(), optimizer_path) |
|
|
|
|
|
def test(): |
|
network.eval() |
|
test_loss = 0 |
|
correct = 0 |
|
with torch.no_grad(): |
|
for data, target in test_loader: |
|
output = network(data) |
|
test_loss += F.nll_loss(output, target, size_average=False).item() |
|
pred = output.data.max(1, keepdim=True)[1] |
|
correct += pred.eq(target.data.view_as(pred)).sum() |
|
test_loss /= len(test_loader.dataset) |
|
test_losses.append(test_loss) |
|
print('\nTest set: Avg. loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format( |
|
test_loss, correct, len(test_loader.dataset), |
|
100. * correct / len(test_loader.dataset))) |
|
|
|
|
|
for epoch in range(1, n_epochs + 1): |
|
train(epoch) |
|
test() |
|
|
|
fig = plt.figure() |
|
plt.plot(train_counter, train_losses, color='blue') |
|
print(len(test_counter)) |
|
print(len(test_losses)) |
|
plt.scatter(test_counter, test_losses, color='red') |
|
plt.legend(['Train Loss', 'Test Loss'], loc='upper right') |
|
plt.xlabel('number of training examples seen') |
|
plt.ylabel('negative log likelihood loss') |
|
plt.show() |
|
|
