PyTorch—— softmax 的从零开始实现

本文是学习《动手学深度学习（pytorch）》“3.6 softmax 的从零开始实现” 的笔记，具体解释请参考原文。

#导入包
import torch
import torchvision
import torch.utils.data as Data

一、获取数据

参考 PyTorch—— 图像分类数据集（Fashion-MNIST），把“一、二”两部分整合起来写成一个函数：

def load_data_fashion_mnist(batch_size):
	mnist_train = torchvision.datasets.FashionMNIST(root='~/Datasets/FashionMNIST', train=True, download=True, transform=torchvision.transforms.ToTensor())
	mnist_test = torchvision.datasets.FashionMNIST(root='~/Datasets/FashionMNIST', train=False, download=True, transform=torchvision.transforms.ToTensor())

	train_iter = Data.DataLoader(dataset=mnist_train, batch_size=batch_size, shuffle=True)
	test_iter = Data.DataLoader(dataset=mnist_test, batch_size=batch_size, shuffle=True)
	return train_iter, test_iter

batch_size = 256
train_iter, test_iter = load_data_fashion_mnist(batch_size)

二、初始化模型参数

需要模型参数梯度

num_inputs = 784
num_outputs = 10

weights = torch.normal(0, 1, (num_inputs, num_outputs))
bias = torch.zeros(num_outputs)

weights.requires_grad_(requires_grad=True)
bias.requires_grad_(requires_grad=True) 

三、实现softmax运算

softmax运算原理介绍

def softmax(X):
	X_exp = X.exp()
	partition = X_exp.sum(dim=1, keepdim=True)
	return X_exp / partition

四、定义模型

def net(X):
	return softmax(torch.mm(X.view(-1, num_inputs), weights) + bias)

五、定义损失函数

损失函数使用交叉熵函数。gather函数的第一个参数是维度。

def cross_entropy(y, y_hat):
	return -torch.log(y_hat.gather(1, y.view(-1,1)))

六、计算分类准确率

1、首先是训练集的分类准确率计算函数，只用计算 train_iter 中当前 batch_size 大小的数据的分类准确率即可。

def accuracy(y, y_hat):
	return (y_hat.argmax(dim=1)==y).float().mean().item()

2、第二个是测试集的分类准确率计算函数，需要计算 test_iter 中所有数据的分类准确率。

def evaluate_accuracy(data_iter, net):
	acc_sum, num = 0.0, 0
	for X, y in data_iter:
		acc_sum += (net(X).argmax(dim=1)==y).float().sum().item()
		num += y.shape[0]
	return acc_sum / num

七、定义优化函数

def sgd(params, batch_size, lr):
	for param in params:
		param.data -= lr * param.grad / batch_size

八、训练模型

num_epochs = 5
lr = 0.5

for epoch in range(num_epochs):
	acc_sum, l_sum, num = 0.0, 0.0, 0
	for X, y in train_iter:
		y_hat = net(X)
		l = cross_entropy(y, y_hat).sum()
		l_sum += l
		l.backward()

		sgd([weights, bias], batch_size, lr)
		weights.grad.data.zero_()
		bias.grad.data.zero_()
		
		acc_sum += accuracy(y, y_hat)
		num += y.shape[0]
	test_acc = evaluate_accuracy(test_iter, net)
	print('Step:%d, Loss:%.3f, train accuracy:%.3f, test accuracy:%.3f' % (epoch+1, l_sum/num, acc_sum/num, test_acc))

我们把训练模型封装成一个函数，由外界决定：

• 优化函数是自定义 还是 使用pytorch中自带的函数
• 损失函数是自定义 还是 使用pytorch中自带的函数

def train_softmax(net, train_iter, test_iter, loss, num_epochs, batch_size, params=None, lr=None, optimizer=None):
	for epoch in range(num_epochs):
		acc_sum, l_sum, num = 0.0, 0.0, 0
		for X, y in train_iter:
			y_hat = net(X)
			l = loss(y, y_hat).sum()
			l_sum += l#.item()
	
			if optimizer is not None:
				optimizer.zero_grad()
			elif params is not None and params[0].grad is not None:
				for param in params:
					param.grad.data.zero_()
			l.backward()
			if optimizer is not None:
				optimizer.step()
			else:
				sgd(params, batch_size, lr)
			
			acc_sum += accuracy(y, y_hat)
			num += y.shape[0]
		test_acc = evaluate_accuracy(test_iter, net)
		print('Step:%d, Loss:%.3f, train accuracy:%.3f, test accuracy:%.3f' % (epoch+1, l_sum/num, acc_sum/num, test_acc))

调用：

train_softmax(net, train_iter, test_iter, cross_entropy, num_epochs, batch_size, params=[weights, bias], lr=lr)

#Step:1, Loss:2.370, train accuracy:0.002, test accuracy:0.713
#Step:2, Loss:1.362, train accuracy:0.003, test accuracy:0.761
#Step:3, Loss:1.164, train accuracy:0.003, test accuracy:0.779
#Step:4, Loss:1.077, train accuracy:0.003, test accuracy:0.789
#Step:5, Loss:1.012, train accuracy:0.003, test accuracy:0.789