代码在colab上测试
import os
import datetime
import torch
import torchvision
from torch import nn
from torch import optim
import torch.nn.functional as F
from torch.autograd import Variable
from torch.utils.data import DataLoader
from torchvision import transforms
import torchvision.models as models
from torchvision.utils import save_image
from torchvision.datasets import MNIST
import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import norm
from PIL import Image
import time
if not os.path.exists('./img'):
os.mkdir('./img')
def try_gpu(i=0):
"""如果存在，则返回gpu(i)，否则返回cpu()"""
if torch.cuda.device_count() >= i + 1:
return torch.device(f'cuda:{i}')
return torch.device('cpu')
img_transform = transforms.Compose([
# transforms.Resize([224, 224]),
transforms.ToTensor(),
# transforms.Lambda(lambda x: x.repeat(3, 1, 1)),
# transforms.Normalize([0.5], [0.5])
])
simple_transform = transforms.Compose([
# transforms.Resize((224,224))
transforms.ToTensor()
# ,transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
dataset= MNIST(root='/content/data', train=True, download=True, transform=simple_transform)
length=len(dataset)
train_size,valid_size,test_size=int(0.7*length),int(0.2*length),int(0.1*length)
train_set,valid_set,test_set=torch.utils.data.random_split(dataset,[train_size,valid_size,test_size])
print(len(train_set),len(valid_set),len(test_set))
dataset_sizes = {
'train': len(train_set),
'valid': len(valid_set),
'test': len(test_set)
}
class LeNet(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(1, 6, kernel_size=5, padding=2)
self.conv2 = nn.Conv2d(6, 16, kernel_size=5)
self.sigmoid = nn.Sigmoid()
self.bn = nn.AvgPool2d(kernel_size=2,stride=2)
self.linear1 = nn.Linear(16 * 5 * 5, 120)
self.linear2 = nn.Linear(120, 84)
self.linear3 = nn.Linear(84, 10)
self.fn = nn.Flatten()
def forward(self, x):
out = self.bn(self.sigmoid(self.conv1(x)))
out = self.bn(self.sigmoid(self.conv2(out)))
out = self.fn(out)
out = self.linear1(out)
out = self.linear2(out)
out = self.linear3(out)
return out
net = nn.Sequential(
nn.Conv2d(1, 6, kernel_size=5, padding=2), nn.Sigmoid(),
nn.AvgPool2d(kernel_size=2, stride=2),
nn.Conv2d(6, 16, kernel_size=5), nn.Sigmoid(),
nn.AvgPool2d(kernel_size=2, stride=2),
nn.Flatten(),
nn.Linear(16 * 5 * 5, 120), nn.Sigmoid(),
nn.Linear(120, 84), nn.Sigmoid(),
nn.Linear(84, 10))
X = torch.rand(size=(1, 1, 28, 28), dtype=torch.float32)
for layer in net:
X = layer(X)
print(layer.__class__.__name__,'output shape: \t',X.shape)
val_losses,val_acc = [],[]
train_losses,train_acc = [],[]
def train_model(model, criterion, optimizer, scheduler, device, num_epochs, dataloaders):
model.to(device)
print('training on ', device)
since = time.time()
best_model_wts = []
best_acc = 0.0
for epoch in range(num_epochs):
for phase in ['train','valid']:
s = time.time()
if phase == 'tarin':
model.train()
else:
model.eval（)
running_loss = 0.0
running_corrects = 0
for data in dataloaders[phase]:
inputs,labels = data
if torch.cuda.is_available():
inputs = Variable(inputs.to(device))
labels = Variable(labels.to(device))
else:
inputs = Variable(inputs)
labels = Variable(labels)
if phase == 'train':
optimizer.zero_grad()
outputs = model(inputs)
_,preds = torch.max(outputs.data,1)
loss = criterion(outputs,labels)
if phase == 'train':
loss.backward()
optimizer.step()
running_loss += loss.data
running_corrects += torch.sum(preds == labels.data)
epoch_loss = running_loss / dataset_sizes[phase]
epoch_acc = running_corrects / dataset_sizes[phase]
if phase == 'valid':
val_losses.append(epoch_loss.to('cpu'))
val_acc.append(epoch_acc.to('cpu'))
if phase == 'train':
train_losses.append(epoch_loss.to('cpu'))
train_acc.append(epoch_acc.to('cpu'))
print('Epoch {}/{} - {} Loss: {:.4f} Acc: {:.4f} Time:{:.1f}s'
.format(epoch+1, num_epochs, phase, epoch_loss, epoch_acc,time.time()-s))
# 每轮都记录最好的参数.
if phase == 'valid' and epoch_acc > best_acc:
best_acc = epoch_acc
best_model_wts = model.state_dict()
if phase == 'train':
scheduler.step()
print()
time_elapsed = time.time() - since
print('Training complete in {:.0f}m {:.0f}s'.format(
time_elapsed // 60, time_elapsed % 60))
print('Best val Acc: {:4f}'.format(best_acc))
# model.load_state_dict(best_model_wts)
return model
batch_size = 128
dataloader_train = DataLoader(train_set, batch_size=batch_size, shuffle=True)
dataloader_valid = DataLoader(valid_set, batch_size=batch_size)
dataloader_test = DataLoader(test_set, batch_size=batch_size)
dataloaders = {
'train': dataloader_train,
'valid': dataloader_valid,
'test': dataloader_test
}
num_epochs = 10
model = LeNet()
criterion = nn.CrossEntropyLoss()
# optimizer = optim.Adam(model.parameters(), lr=0.1)
optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)
model = train_model(model, criterion, optimizer, scheduler, try_gpu(), num_epochs, dataloaders)
plt.plot(range(1, len(train_losses)+1),train_losses, 'b', label='training loss')
plt.plot(range(1, len(val_losses)+1), val_losses, 'r', label='val loss')
plt.legend()
plt.plot(range(1,len(train_acc)+1),train_acc,'b--',label = 'train accuracy')
plt.plot(range(1,len(val_acc)+1),val_acc,'r--',label = 'val accuracy')
plt.legend()
def test_model(model, criterion, device, dataloaders):
model.to(device)
print('Testing on ', device)
s = time.time()
model.eval（)
running_loss = 0.0
running_corrects = 0
phase = "test"
for inputs,labels in dataloaders[phase]:
if torch.cuda.is_available():
inputs = Variable(inputs.to(device))
labels = Variable(labels.to(device))
else:
inputs = Variable(inputs)
labels = Variable(labels)
outputs = model(inputs)
_,preds = torch.max(outputs.data,1)
loss = criterion(outputs,labels)
running_loss += loss.data
running_corrects += torch.sum(preds == labels.data)
epoch_loss = running_loss / dataset_sizes[phase]
epoch_acc = running_corrects / dataset_sizes[phase]
print('Loss: {:.4f} Acc: {:.4f} Time:{:.1f}s'
.format(epoch_loss, epoch_acc,time.time()-s))
# 显示图片
def imshow(inp):
inp = inp.numpy().transpose((1, 2, 0))
mean = np.array([0.485, 0.456, 0.406])
std = np.array([0.229, 0.224, 0.225])
inp = std * inp + mean
inp = np.clip(inp, 0, 1)
plt.imshow(inp)
imshow(dataset[1][0])
torch.save(model, '/content/lenet.pth')
model_dict=torch.load('/content/lenet.pth')
# print(model_dict)
for i, (x, y) in enumerate(dataloaders['test']):
if i == 0:
x = x[0].resize(1,1,28,28)
outputs = model_dict( Variable( x.to(try_gpu(0)) ) )
_,preds = torch.max(outputs.data,1)
print(preds)
print(y[0])
