1、导入库
from torchvision import transforms
from torch.utils import data
import matplotlib.pyplot as plt
import torch
import torchvision
import seaborn as sns
import torch.nn as nn
import numpy as np
2、工具函数
2.1 计数器
class Accumulator():
def __init__(self, n):
self.data = [0.0]*n
def add(self,*args):
self.data = [a + float(b) for a,b in zip(self.data, args)]
def reset(self):
self.data = [0.0]*len(self.data)
def __getitem__(self, idx):
return self.data[idx]
2.2 图片显示函数
def show_images(imgs, labels, num_rows, num_cols, title=None, scale=1.5):
sns.set(font_scale=0.9, style='whitegrid')
figsize = (num_rows*scale, num_cols*scale)
fig = plt.figure(figsize=figsize)
for i,(img,label) in enumerate(zip(imgs, labels)):
ax = fig.add_subplot(num_rows, num_cols, i+1)
ax.set_title(f'label = {label}')
plt.imshow(img.reshape(28,28).numpy())
plt.tight_layout()
plt.show()
2.3 损失函数绘图
def draw_loss(train_loss, test_loss):
n = len(train_loss)
x = np.linspace(1,n,n)
fig, ax = plt.subplots()
ax.plot(x,train_loss,label='train_loss')
ax.plot(x,test_loss,label='test_loss')
ax.set_xlabel('epoch')
ax.set_ylabel('loss')
ax.set_title('img_calssfier')
ax.legend()
plt.show()
2.4 计算分类正确的数量
def accuracy(y_hat, y):
if len(y_hat.shape)>1 and y_hat.shape[1]>1 :
y_hat = y_hat.argmax(axis=1)
cmp = y_hat.type(y.dtype) == y
return float(cmp.type(y.dtype).sum())
3.数据集处理
def load_data_fashion_mnist(batch_size, resize=None):
trans = [transforms.ToTensor()]
if resize:
trans.insert(0, transforms.Resize(resize))
trans = transforms.Compose(trans)
mnist_train = torchvision.datasets.FashionMNIST(
root='../data/', train=True, transform=trans, download=True)
mnist_test = torchvision.datasets.FashionMNIST(
root='../data/', train=False, transform=trans, download=True)
return (data.DataLoader(mnist_train, batch_size=batch_size, shuffle=True,num_workers=4),
data.DataLoader(mnist_test, batch_size=batch_size, shuffle=False,num_workers=4))
4. 定义网络模型
net = nn.Sequential(nn.Flatten(), nn.Linear(28*28, 10))
5. 定义优化器
trainer = torch.optim.SGD(net.parameters(), lr=0.1)
6. 定义损失函数
loss = nn.CrossEntropyLoss(reduction='none')
7. 训练过程
训练一个epoch
def epoch_train(net, train_iter,loss, optim):
#训练损失总和、正确总和、样本总和
metric = Accumulator(3)
net.train()
for x,y in train_iter:
l = loss(net(x),y)
optim.zero_grad()
l.mean().backward()
optim.step()
metric.add(float(l.sum()), accuracy(net(x),y), y.numel())
return metric[0]/metric[2], metric[1]/metric[2]
测试函数
def test(test_iter,net,loss):
net.eval()
accumulator = Accumulator(3)
with torch.no_grad():
for x,y in test_iter:
y_hat = net(x)
l = loss(y_hat, y)
accumulator.add(float(l.sum()), accuracy(y_hat,y), y.numel())
return accumulator[0]/accumulator[2], accumulator[1]/accumulator[2]
训练函数
def train(num_epochs, train_iter, test_iter, net,optim,loss):
loss_1,loss_2 = [],[]
for epoch in range(num_epochs):
train_loss, train_accuracy = epoch_train(net,train_iter=train_iter, optim=optim,loss=loss)
#模型测试
test_loss, test_accuracy = test(test_iter=test_iter, net=net, loss=loss)
loss_1.append(train_loss)
loss_2.append(test_loss)
print(f'epoch:{epoch+1}, train_loss:{train_loss},train_accuracy:{train_accuracy}')
draw_loss(loss_1, loss_2)
开始训练
num_epochs=10
train_iter, test_iter = load_data_fashion_mnist(batch_size=100)
train(num_epochs=num_epochs, train_iter=train_iter,test_iter=test_iter,net=net, optim=trainer,loss=loss)