NanoDL: 基于Jax的轻量级深度学习库

NanoDL: 为深度学习爱好者打造的轻量级工具箱

在人工智能和深度学习迅速发展的今天,Transformer模型已经成为许多自然语言处理和计算机视觉任务的首选架构。然而,设计和训练这些复杂的模型往往需要大量的计算资源和专业知识。为了让更多的研究者和开发者能够轻松地探索和创新Transformer模型,Henry Ndubuaku开发了NanoDL这个基于Jax的轻量级深度学习库。

NanoDL的核心特性

NanoDL的设计理念是"小而美",它专注于提供一套简洁而强大的工具,让用户可以从头开始设计和训练Transformer模型。以下是NanoDL的一些核心特性:

丰富的模块和层: NanoDL提供了广泛的构建块和层,使用户可以自由组合创建定制的Transformer模型。
流行模型实现: 库中包含了多种流行模型的实现,如Gemma、LlaMa3、Mistral、GPT3、GPT4(推断)、T5、Whisper、ViT、Mixers和CLIP等。
分布式训练支持: NanoDL集成了数据并行的分布式训练器,可以在多个GPU或TPU上进行模型训练,无需手动编写训练循环。
高效的数据加载: 提供了专门为Jax/Flax设计的数据加载器,简化了数据处理流程。
独特的注意力机制: 实现了一些在Flax/Jax中不常见的层,如RoPE、GQA、MQA和SWin注意力等。
GPU/TPU加速的经典机器学习模型: 包括PCA、KMeans、回归和高斯过程等。
真随机数生成器: 简化了Jax中随机数生成的复杂过程。
NLP和计算机视觉算法: 实现了高斯模糊、BLEU分数计算、分词器等多种高级算法。
模块化设计: 每个模型都被封装在单个文件中,没有外部依赖,方便用户直接使用或修改源代码。

NanoDL Logo

快速上手NanoDL

要开始使用NanoDL,你首先需要安装Python 3.9或更高版本,以及JAX、FLAX和OPTAX库。对于只需要CPU版本的用户,可以通过以下命令安装:

pip install --upgrade pip
pip install jax flax optax

然后,通过PyPI安装NanoDL:

pip install nanodl

NanoDL的实际应用案例

为了更好地理解NanoDL的使用方法,让我们来看几个具体的应用案例:

1. 文本生成模型

以下是使用NanoDL训练一个简单的GPT4模型的示例代码:

import jax
import nanodl
import jax.numpy as jnp
from nanodl import ArrayDataset, DataLoader
from nanodl import GPT4, GPTDataParallelTrainer

# 准备数据集
batch_size = 8
max_length = 50
vocab_size = 1000

# 创建随机数据
data = nanodl.uniform(
    shape=(batch_size, max_length), 
    minval=0, maxval=vocab_size-1
    ).astype(jnp.int32)

# 创建输入和目标
dummy_inputs, dummy_targets = data[:, :-1], data[:, 1:]

# 创建数据集和数据加载器
dataset = ArrayDataset(dummy_inputs, dummy_targets)
dataloader = DataLoader(
    dataset, batch_size=batch_size, shuffle=True, drop_last=False
    )

# 模型参数
hyperparams = {
    'num_layers': 1,
    'hidden_dim': 256,
    'num_heads': 2,
    'feedforward_dim': 256,
    'dropout': 0.1,
    'vocab_size': vocab_size,
    'embed_dim': 256,
    'max_length': max_length,
    'start_token': 0,
    'end_token': 50,
}

# 初始化GPT4模型
model = GPT4(**hyperparams)

# 创建训练器
trainer = GPTDataParallelTrainer(
    model, dummy_inputs.shape, 'params.pkl'
    )

# 开始训练
trainer.train(
    train_loader=dataloader, num_epochs=100, val_loader=dataloader
    )

# 加载训练好的参数
params = trainer.load_params('params.pkl')

# 生成文本
start_tokens = jnp.array([[123, 456]])
outputs = model.apply(
    {'params': params}, 
    start_tokens,
    rngs={'dropout': nanodl.time_rng_key()}, 
    method=model.generate
    )

这个例子展示了如何使用NanoDL创建一个简单的GPT4模型,训练它,然后用它来生成文本。

2. 图像生成模型

NanoDL也支持图像处理任务。以下是使用扩散模型生成图像的示例:

import nanodl
import jax.numpy as jnp
from nanodl import ArrayDataset, DataLoader
from nanodl import DiffusionModel, DiffusionDataParallelTrainer

image_size = 32
block_depth = 2
batch_size = 8
widths = [32, 64, 128]
input_shape = (101, image_size, image_size, 3)
images = nanodl.normal(shape=input_shape)

# 创建数据集和数据加载器
dataset = ArrayDataset(images) 
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, drop_last=False) 

# 创建扩散模型
diffusion_model = DiffusionModel(image_size, widths, block_depth)

# 训练模型
trainer = DiffusionDataParallelTrainer(diffusion_model, 
                                       input_shape=images.shape, 
                                       weights_filename='params.pkl', 
                                       learning_rate=1e-4)

trainer.train(dataloader, 10)

# 生成图像样本
params = trainer.load_params('params.pkl')
generated_images = diffusion_model.apply({'params': params}, 
                                         num_images=5, 
                                         diffusion_steps=5, 
                                         method=diffusion_model.generate)