赞
踩
import torch
import matplotlib.pyplot as plt
import os
#data
import numpy as np
import torchvision.datasets as dataset
import torchvision.transforms as transforms
import torch.utils.data as data_utils
train_data = dataset.MNIST(root="mnist",
train=True,
transform=transforms.ToTensor(),
download=True)
test_data = dataset.MNIST(root="mnist",
train=False,
transform=transforms.ToTensor(),
download=False)
#batchsize
train_loader =data_utils.DataLoader(dataset = train_data,batch_size = 64,shuffle=True)
test_loader = data_utils.DataLoader(dataset = test_data,batch_size = 64,shuffle=True)
#net
class CNN(torch.nn.Module):
def __init__(self):
super(CNN, self).__init__()
self.conv = torch.nn.Sequential(
torch.nn.Conv2d(1,32,kernel_size=5,padding=2),
torch.nn.BatchNorm2d(32),
torch.nn.ReLU(),
torch.nn.MaxPool2d(2)
)
self.fc = torch.nn.Linear(14*14*32,10)
def forward(self,x):
out = self.conv(x)
out = out.view(out.size()[0],-1)
out = self.fc(out)
return out
cnn=CNN()
cnn = cnn.cuda()
#loss
loss_func = torch.nn.CrossEntropyLoss()
#optimizer
optimizer = torch.optim.Adam(cnn.parameters(),lr=0.01)
#train
for epoch in range(10):
for i,(images,label)in enumerate(train_loader):
images = images.cuda()
labels =label.cuda()
outputs = cnn(images)
loss = loss_func(outputs ,labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
print("eporch is {},{}/{},loss is {}".format(epoch+1,i,len(train_data)//64,loss.item()))
#eval/test
loss_test = 0
loss1 = []
acc1 = []
acc = 0
for i,(images,label)in enumerate(test_loader):
images = images.cuda()
labels = label.cuda()
outputs = cnn(images)
loss_test += loss_func(outputs, labels)
_,pred = outputs.max(1)
acc += (pred ==labels).sum().item()
loss_test = loss_test/(len(test_data)//64)
acc = acc/len(test_data)
print("eporch is {},acc is {},loss is {}".format(epoch+1,acc,loss_test))
torch.save(cnn,"model/model.pkl")
#load
import torch
import torchvision.datasets as dataset
import torchvision.transforms as transforms
import torch.utils.data as data_utils
class CNN(torch.nn.Module):
def __init__(self):
super(CNN, self).__init__()
self.conv = torch.nn.Sequential(
torch.nn.Conv2d(1,32,kernel_size=5,padding=2),
torch.nn.BatchNorm2d(32),
torch.nn.ReLU(),
torch.nn.MaxPool2d(2)
)
self.fc = torch.nn.Linear(14*14*32,10)
def forward(self,x):
out = self.conv(x)
out = out.view(out.size()[0],-1)
out = self.fc(out)
return out
test_data = dataset.MNIST(root="mnist",
train=False,
transform=transforms.ToTensor(),
download=False)
#batchsize
test_loader = data_utils.DataLoader(dataset = test_data,batch_size = 64,shuffle=True)
#net
cnn = torch.load('model/model.pkl')
cnn = cnn.cuda()
#eval/test
loss_test = 0
acc = 0
for i,(images,label)in enumerate(test_loader):
images = images.cuda()
labels = label.cuda()
outputs = cnn(images)
_,pred = outputs.max(1)
acc += (pred ==labels).sum().item()
acc = acc/len(test_data)
print("acc is {}".format(acc)
)
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。