1. 建立LeNet5主干网络
import torch
from torch import nn
# 定义网络模型
class LeNet5(nn.Module):
#初始化网络
def __init__(self):
super(LeNet5, self).__init__()
self.c1 = nn.Conv2d(in_channels=1, out_channels=64, kernel_size=5, padding=2)
self.Sigmoid = nn.Sigmoid()
self.s2 = nn.AvgPool2d(kernel_size=2, stride=2)
self.c3 = nn.Conv2d(in_channels=64, out_channels=128, kernel_size=5)
self.s4 = nn.AvgPool2d(kernel_size=2, stride=2)
self.c5 = nn.Conv2d(in_channels=128, out_channels=256, kernel_size=5)
self.flatten = nn.Flatten()
self.f6 = nn.Linear(in_features=256, out_features=512)
self.output = nn.Linear(in_features=512, out_features=10)
def forward(self, x):
x = self.Sigmoid(self.c1(x))
x = self.s2(x)
x = self.Sigmoid(self.c3(x))
x = self.s4(x)
x = self.c5(x)
x = self.flatten(x)
x = self.f6(x)
x = self.output(x)
return x
if __name__=="__main__":
x = torch.rand([1, 1, 28, 28])
model = LeNet5()
y = model(x)2. 加载MNIST数据集，训练网络，保存最优网络模型
import torch
from torch import nn
from net import LeNet5
from torch.optim import lr_scheduler
from torchvision import datasets, transforms
import os
#数据转化为tensor格式
data_transform = transforms.Compose([
transforms.ToTensor()
])
#加载训练数据集
train_dataset = datasets.MNIST(root='./dataset', train=True, transform=data_transform, download=True)
train_dataloader = torch.utils.data.DataLoader(dataset = train_dataset, batch_size=16, shuffle=True)
#加载测试数据集
test_dataset = datasets.MNIST(root='./dataset', train=False, transform=data_transform, download=True)
test_dataloader = torch.utils.data.DataLoader(dataset = test_dataset, batch_size=16, shuffle=True)
#如果有显卡可以转到GPU
device = 'cuda' if torch.cuda.is_available() else 'cpu'
#调用net里定义的网络模型LeNet5 将数据模型转到GPU
model = LeNet5().to(device)
#定义损失函数
loss_fn = nn.CrossEntropyLoss()
#定义优化器
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3, momentum=0.9)
#学习率每隔10轮变为原来的0.1
lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)
#定义训练函数
def train(dataloader, model, loss_fn, optimizer):
loss, current, n = 0.0, 0.0, 0
for batch, (X, y) in enumerate(dataloader):
X, y = X.to(device), y.to(device)
output = model(X)
cur_loss = loss_fn(output,y)
_, pred = torch.max(output, axis=1)
cur_acc = torch.sum(y == pred)/output.shape[0]
optimizer.zero_grad()
cur_loss.backward()
optimizer.step()
loss += cur_loss.item()
current += cur_acc.item()
n = n + 1
print("train_loss " + str(loss/n))
print("train_acc " + str(current/n))
def val(dataloader, model, loss_fn):
model.eval()
loss, current, n = 0.0, 0.0, 0
with torch.no_grad():
for batch, (X, y) in enumerate(dataloader):
X, y = X.to(device), y.to(device)
output = model(X)
cur_loss = loss_fn(output,y)
_, pred = torch.max(output, axis=1)
cur_acc = torch.sum(y == pred)/output.shape[0]
loss += cur_loss.item()
current += cur_acc.item()
n = n + 1
print("test_loss " + str(loss/n))
print("test_acc " + str(current/n))
return current/n
#开始训练
epoch = 200
min_acc = 0
for t in range(epoch):
print(f'epoch{t+1}\n-----------------------------------')
train(train_dataloader, model, loss_fn, optimizer)
a = val(test_dataloader, model, loss_fn)
#保存最优模型
if a > min_acc:
folder = 'save_model'
if not os.path.exists(folder):
os.mkdir('save_model')
min_acc = a
print('save best model!!!!!!')
torch.save(model.state_dict(), 'save_model/best_model.pth')
print('OK!!!!!!!!!!!!!!!!!')3. 训练过程
epoch1
-----------------------------------
train_loss ?2.30177958946228
train_acc ?0.11585
test_loss ?2.2867366130828857
test_acc ?0.1356
save best model!!!!!!
4. 模型测试
import torch
from net import LeNet5
from torch.autograd import Variable
from torchvision import datasets, transforms
from torchvision.transforms import ToPILImage
#数据转化为tensor格式
data_transform = transforms.Compose([
transforms.ToTensor()
])
#加载训练数据集
train_dataset = datasets.MNIST(root='./dataset', train=True, transform=data_transform, download=True)
train_dataloader = torch.utils.data.DataLoader(dataset = train_dataset, batch_size=16, shuffle=True)
#加载测试数据集
test_dataset = datasets.MNIST(root='./dataset', train=False, transform=data_transform, download=True)
test_dataloader = torch.utils.data.DataLoader(dataset = test_dataset, batch_size=16, shuffle=True)
#如果有显卡可以转到GPU
device = 'cuda' if torch.cuda.is_available() else 'cpu'
#调用net里定义的网络模型LeNet5 将数据模型转到GPU
model = LeNet5().to(device)
model.load_state_dict(torch.load(r"C:\scj\research\pytorch\save_model\best_model.pth"))
#获取结果
classes = [
"0",
"1",
"2",
"3",
"4",
"5",
"6",
"7",
"8",
"9",
]
#把tensor转化为图片，方便可视化
show = ToPILImage()
#验证
print(len(test_dataset))
for i in range(len(test_dataset)):
X, y = test_dataset[i][0], test_dataset[i][1]
# show(X).show()
X = Variable(torch.unsqueeze(X, dim=0).float(), requires_grad=False).to(device)
with torch.no_grad():
pred = model(X)
predicted, actual = classes[torch.argmax(pred[0])], classes[y]
print(f'figure {str(i)}: predicted: "{predicted}", actual: "{actual}"')
