使用 PyTorch 数据读取,JAX 框架来训练一个简单的神经网络

使用 PyTorch 数据读取,JAX 框架来训练一个简单的神经网络

本文例程部分主要参考官方文档。

JAX简介

JAX 的前身是 Autograd ,也就是说 JAX 是 Autograd 升级版本,JAX 可以对 Python 和 NumPy 程序进行自动微分。可以通过 Python的大量特征子集进行区分,包括循环、分支、递归和闭包语句进行自动求导,也可以求三阶导数(三阶导数是由原函数导数的导数的导数。 所谓三阶导数,即原函数导数的导数的导数,将原函数进行三次求导)。通过 grad ,JAX 支持反向模式和正向模式的求导,而且这两种模式可以任意组合成任何顺序,具有一定灵活性。

另一个特点是基于 XLA 的 JIT 即时编译,大大提高速度。

需要注意的是,JAX 仅提供计算时的优化,相当于是一个支持自动微分和 JIT 编译的 NumPy。也就是说,数据处理 Dataloader 等其他框架都会提供的 utils 功能这里是没有的。所幸 JAX 可以比较好的支持 PyTorch、 TensorFlow 等主流框架的数据读取。本文就将基于 PyTorch 的数据读取工具和 JAX 框架来训练一个简单的神经网络

以下是国内优秀的机器学习框架 OneFlow 同名公司的创始人袁进辉老师在知乎上的一个评价:

如果说tensorflow 是主打lazy, 偏functional 的思想,但实现的臃肿面目可憎;pytorch 则主打eager, 偏imperative 编程,但内核简单,可视为支持gpu的numpy, 加上一个autograd。JAX 像是这俩框架的混合体,取了tensorflow的functional和PyTorch的精简,即支持gpu的 numpy, 具有autograd功能,非常追求函数式编程的思想,强调无状态,immutable,加上JIT修饰符后就是lazy,可以使用xla对计算流程进行静态分析和优化。当然JAX不带jit也可像pytorch那种命令式编程和eager执行。

JAX有可能和PyTorch竞争。

安装

安装可以通过源码编译,也可以直接 pip。源码编译详见[官方文档: Building from source][2],对于官方没有提供预编译包的 cuda-cudnn 版本组合,只能通过自己源码构建。pip的方式比较简单,在 github 仓库的 README 文档中就有介绍。要注意,不同于 PyTorch 等框架,JAX 不会再 pip 安装中绑定 CUDA 或 cuDNN 进行安装,若未安装,需要自己先手动安装。仅使用 CPU 的版本也有支持。

笔者是 CUDA11.1,CUDNN 8.2,安装如下:

pip install --upgrade pip
pip install jax[cuda11_cudnn82] -f https://storage.googleapis.com/jax-releases/jax_releases.html

前面已经提到过,本文会借用 PyTorch 的数据处理工具,因此 torch 和 torchvision 也是必不可少的(已经安装的可跳过):

pip install torch torchvision

构建简单的神经网络训练

框架安装完毕,我们正式开始。接下来我们使用 JAX 在 MNIST 上指定和训练一个简单的 MLP 进行计算,用 PyTorch 的数据加载 API 来加载图像和标签。

import jax.numpy as jnp
from jax import grad, jit, vmap
from jax import random

超参数

# 本函数用来随机初始化网络权重
def random_layer_params(m, n, key, scale=1e-2):
	w_key, b_key = random.split(key)
	return scale * random.normal(w_key, (n, m)), scale * random.normal(b_key, (n, ))

# 初始化各个全连接层
def init_network_params(sizes, key):
	keys = random.split(key, len(sizes))
	return [random_layer_params(m, n, k) for m, n, k in zip(sizes[: -1], sizes[1: ], keys)]

layer_sizes = [784, 512, 512, 10]
step_size = 0.01
num_epochs = 8
batch_size = 128
n_targets = 10
params = init_network_params(layer_sizes, random.PRNGKey(0))

自动分批次预测

对于小批量,我们稍后将使用 JAX 的 vmap 函数来自动处理,而不会降低性能。我们现在先准备一个单张图像推理预测函数:

from jax.scipy.special import logsumexp

def relu(x):
	return jnp.maximum(0, x)

# 对单张图像进行推理的函数
def predict(params, image):
	activations = image
	for w, b in params[: -1]:
		outputs = jnp.dot(w, activations) + b
		activations = relu(outputs)
	
	final_w, final_b = params[-1]
	logits = jnp.dot(final_w, activations) + final_b
	return logits - logsumexp(logits)

这个函数应该只能用来处理单张图像推理预测,而不能批量处理,我们简单测试一下,对于单张:

random_flattened_images = random.normal(random.PRNGKey(1), (28 * 28,))
preds = predict(params, random_flattened_images)
print(preds.shape)

输出:

(10,)

对于批次:

random_flattened_images = random.normal(random.PRNGKey(1), (10, 28 * 28))
try:
	preds = predict(params, random_flattened_images)
except TypeError:
	print('Invalid shapes!')

输出:

Invalid shapes!

现在我们使用 vmap 来使它能够处理批量数据:

# 用 vmap 来实现一个批量版本
batched_predict = vmap(predict, in_axes=(None, 0))

# batched_predict 的调用与 predict 相同
batched_preds = batched_predict(params, random_flattened_images)
print(batched_preds.shape)

输出:

(10, 10)

现在,我们已经做好了准备工作,接下来就是要定义一个神经网络并且进行训练了,我们已经构建了的自动批处理版本的 predict 函数,并且将在损失函数中也使用它。我们将使用 grad 来得到损失关于神经网络参数的导数。而且,这一切都可以用 jit 进行加速。

实用工具函数和损失函数

def one_hot(x, k, dtype=jnp.float32):
	"""构建一个 x 的 k 维 one-hot 编码."""
	return jnp.array(x[:, None] == jnp.arange(k), dtype)


def accuracy(params, images, targets):
	target_class = jnp.argmax(targets, axis=1)
	predicted_class =  jnp.argmax(batched_predict(params, images), axis=1)
	return jnp.mean(predicted_class == target_class)

def loss(params, images, targets):
	preds = batched_predict(params, images)
	return -jnp.mean(preds * targets)

@jit
def update(params, x, y):
	grads = grad(loss)(params, x, y)
	return [(w - step_size * dw, b - step_size * db) for (w, b), (dw, db) in zip(params, grads)]

使用 PyTorch 进行数据读取

JAX 是一个专注于程序转换和支持加速的 NumPy,对于数据的读取,已经有很多优秀的工具了,这里我们就直接用 PyTorch 的 API。我们会做一个小的 shim 来使得它能够支持 NumPy 数组。

import numpy as np
from torch.utils import data
from torchvision.datasets import MNIST

def numpy_collate(batch):
	if isinstance(batch[0], np.ndarray):
		return np.stack(batch)
	elif isinstance(batch[0], (tuple, list)):
		transposed = zip(*batch)
		return [numpy_collate(samples) for samples in transposed]
	else:
		return np.array(batch)

class NumpyLoader(data.DataLoader):
	def __init__(self, dataset, batch_size=1,
					shuffle=False, sampler=None,
					batch_sampler=None, num_workers=0,
					pin_memory=False, drop_last=False,
					timeout=0, worker_init_fn=None):
		super(self.__class__, self).__init__(dataset,
					batch_size=batch_size,
					shuffle=shuffle,
					sampler=sampler,
					batch_sampler=batch_sampler,
					collate_fn=numpy_collate,
					num_workers=num_workers,
					pin_memory=pin_memory,
					drop_last=drop_last,
					timeout=timeout,
					worker_init_fn=worker_init_fn)

class FlattenAndCast(object):
	def __call__(self, pic):
		return np.ravel(np.array(pic, dtype=jnp.float32))

接下来借助 PyTorch 的 datasets,定义我们自己的 dataset:

mnist_dataset = MNIST('/tmp/mnist/', download=True, transform=FlattenAndCast())
training_generator = NumpyLoader(mnist_dataset, batch_size=batch_size, num_workers=0)

此处应该输出一堆下载 MNIST 数据集的信息,就不贴了。

接下来分别拿到整个训练集和整个测试集,下面会用于测准确率:

train_images = np.array(mnist_dataset.train_data).reshape(len(mnist_dataset.train_data), -1)
train_labels = one_hot(np.array(mnist_dataset.train_labels), n_targets)

mnist_dataset_test = MNIST('/tmp/mnist/', download=True, train=False)
test_images = jnp.array(mnist_dataset_test.test_data.numpy().reshape(len(mnist_dataset_test.test_data), -1), dtype=jnp.float32)
test_labels = one_hot(np.array(mnist_dataset_test.test_labels), n_targets)

开始训练

import time
for epoch in range(num_epochs):
	start_time = time.time()
	for x, y in training_generator:
		y = one_hot(y, n_targets)
		params = update(params, x, y)
	epoch_time = time.time() - start_time

	train_acc = accuracy(params, train_images, train_labels)
	test_acc = accuracy(params, test_images, test_labels)
	print("Epoch {} in {:0.2f} sec".format(epoch, epoch_time))
	print("Training set accuracy {}".format(train_acc))
	print("Test set accuracy {}".format(test_acc))

输出:

Epoch 0 in 3.29 sec
Training set accuracy 0.9156666994094849
Test set accuracy 0.9196999669075012
...
Epoch 7 in 1.78 sec
Training set accuracy 0.9736666679382324
Test set accuracy 0.9670999646186829

在本文的过程中,我们已经使用了整个 JAX API:grad 用于自动微分、jit 用于加速、vmap 用于自动矢量化。我们使用 NumPy 来进行我们所有的计算,并从 PyTorch 借用了出色的数据加载器,并在 GPU 上运行了整个过程。

Ref:

https://juejin.cn/post/6994695537316331556

https://jax.readthedocs.io/en/latest/notebooks/Neural_Network_and_Data_Loading.html

https://jax.readthedocs.io/en/latest/developer.html#building-from-source

你可能感兴趣的:(Python,pytorch,神经网络,python)