nanodl
1.0.0
作者：Henry Ndubuaku（Discord 和文件徽章可點擊）
注意：代碼在教學上的實施是以重複為代價的。每個模型都有目的地包含在一個檔案中，沒有檔案間依賴關係。
開發和訓練基於 Transformer 的模型通常是資源密集型且耗時的，AI/ML 專家經常需要針對特定問題建立這些模型的較小規模版本。 Jax 是一個資源匱乏但功能強大的框架，它加速了神經網路的開發並抽象化了分散式訓練，但 Jax 中用於 Transformer 開發的現有資源有限。 NanoDL 透過以下功能應對這項挑戰：
各種各樣的塊和層，有助於從頭開始創建定制的變壓器模型。
廣泛的模型選擇，如 Gemma、LlaMa3、Mistral、GPT3、GPT4（推論）、T5、Whisper、ViT、Mixers、CLIP 等。
資料平行分散式訓練器在多個 GPU 或 TPU 上建模，無需手動訓練循環。
資料載入器，使 Jax/Flax 的資料處理過程更加簡單有效。
Flax/Jax 中沒有的層，例如 RoPE、GQA、MQA 和 SWin 注意力，允許更靈活的模型開發。
GPU/TPU 加速的經典 ML 模型，如 PCA、KMeans、回歸、高斯過程等。
Jax 中的真正隨機數產生器不需要詳細的程式碼。
一系列用於 NLP 和電腦視覺任務的高階演算法，例如高斯模糊、BLEU、Tokenizer 等。
每個模型都包含在一個檔案中，沒有外部依賴關係，因此原始程式碼也可以輕鬆使用。
Jax 中的真正隨機數產生器不需要詳細程式碼（下一節中顯示的範例）。
儲存庫中存在實驗性和/或未完成的功能（例如 MAMBA、KAN、BitNet、GAT 和 RLHF），這些功能尚未透過軟體包提供，但可以從此儲存庫複製。歡迎對我們的任何討論、問題和拉取請求線程提供反饋！請在 Discord 中回報任何功能請求、問題、問題或疑慮，或只是讓我們知道您正在做什麼！
您將需要Python 3.9或更高版本，以及可運行的JAX安裝、FLAX安裝、OPTAX安裝（有GPU支援運行訓練，沒有則只能支援創作）。模型可以在 CPU 上設計和測試，但訓練器都是分散式資料並行的，這需要具有 1 到 N 個 GPUS/TPUS 的 GPU。對於僅 CPU 版本的 JAX：
pip install --upgrade pip # To support manylinux2010 wheels. pip install jax flax optax
然後，從 PyPi 安裝 nanodl：
pip install nanodl
我們提供 nanodl API 的各種範例用法。
。 , 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,'hdenden_25. num_heads': 2,'feedforward_dim': 256,' dropout': 0.1,'vocab_size': vocab_size,'embed_dim': 256,'max_length': max_length,'start_token: 0,'
}# 推斷的GPT4 模型model = GPT4(**hyperparams)trainer = GPTDataParallelTrainer(model, dummy_inputs.shape, 'params.pkl')trainer.train(train_loader=dataloader, num_epochs=100, val_loader=100, val= val 資料# 從頭開始產生 tokenstart_tokens = jnp.array([[123, 456]])# 記得載入訓練好的參數 params = trainer.load_params('params.pkl')outputs = model.apply(
{'params': params}, start_tokens,rngs={'dropout': nanodl.time_rng_key()}, method=model.generate)願景範例
導入nanodlimport jax.numpy為jnpfrom nanodl導入ArrayDataset，DataLoaderfrom nanodl導入DiffusionModel，DiffusionDataParallelTrainerimage_size = 32block_深度= 2batch_size = 8widths = [32,64,128]（inputimages） normal（形狀= input_shape)# 使用你自己的圖像資料集= ArrayDataset(images)
dataloader = DataLoader(資料集,batch_size=batch_size,shuffle=True,drop_last=False)
# 建立擴散模型diffusion_model = DiffusionModel(image_size, widths, block_depth)# 在資料訓練器上進行訓練 = DiffusionDataParallelTrainer(diffusion_model,
輸入形狀=圖像.形狀，
weights_filename='params.pkl',
Learning_rate=1e-4)trainer.train(dataloader, 10)# 產生一些樣本：每個模型都是一個Flax.linen 模組# 像平常一樣使用params = trainer.load_params('params.pkl') generated_images = iteration_model. apply( {'參數'：參數}，
圖片數量=5,
擴散步數=5，
方法=diffusion_model.generate）音訊範例
導入jaximport jax.numpy作為nanodl導入陣列dataSet，nanodl導入的dataLoader，whisperdataparallettrainer＃dummy data data parametersbatch_size = 8max_length = 8max_length 101，max_length ), d5,003(Femi_im. _dim) )dataset = ArrayDataset(dummy_inputs, dummy_targets)dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, drop_last=False)# 模型參數hyperparams = {'num_layers'：1，'hidden_s. '：2，'feedforward_dim'：256，'dropout'：0.1，'vocab_size'：1000，'embed_dim'：embed_dim，'max_length'：max_length ,'start_token': 0,'end_token': 50,
}# 初始化模型 model = Whisper(**hyperparams)# 在 datatrainer 上進行訓練 = WhisperDataParallelTrainer(model,
虛擬輸入.形狀，
虛擬目標.形狀，
'params.pkl')trainer.train(dataloader, 2, dataloader)# 樣本推理params = trainer.load_params('params.pkl')# 對於多個樣本，經常使用 model.generate_batchtranscripts = model.apply({'params ': 參數},
dummy_inputs[:1], method=model.generate)RLHF 獎勵模型範例
import nanodlimport jax.numpy as jnpfrom nanodl import ArrayDataset, DataLoaderfrom nanodl import Mistral, RewardModel, RewardDataParallelTrainer# 產生虛擬資料batch_size = 8max_length = 10# 替換為實際標記化資料.int32) dummy_rejected = jnp.zeros((101, max_length), dtype=jnp.int32)# 建立資料集與dataloaderdataset = ArrayDataset(dummy_chosen, dummy_rejected)dataloader = DataLoader(dataset, bone_chosen, dummy_rejected)dataloader = DataLoader(dataset, bize_dsizess,dsize, sh片_pizes)_drm)_sizes)ds)ds), n; False) # 模型參數hyperparams = {'num_layers'：1，'hidden_dim'：256，'num_heads'：2，'feedforward_dim'：256，'dropout'：0.1，'vocab_size'：1000，'embed_dim' ：256，'max_length'：max_length ,'start_token': 0,'end_token': 50,'num_groups': 2,'window_size': 5,'shift_size': 2}# 從Mistralmodel = Mistral(rehyperdelms) = RewardModel(model, dim=hyperparams['hidden_dim'], dropout=0.1)# 訓練獎勵模型trainer = RewardDataParallelTrainer(reward_model, dummy_chosen.shape, 'reward_model_weights.reward_model, dummy_chosen.shape, 'reward_model_weights.reward_model, dummy_chosen.shape, 'reward_model_weights.p. )params = trainer.load_params(' reward_model_weights.pkl')# 像常規 Flax 模型一樣呼叫rewards =reward_model.apply({'params': params}, dummy_chosen, rngs={'dropout': nanodl.time_rng_key(time)}主成分分析範例
import nanodlfrom nanodl import PCA# 使用實際資料data = nanodl.normal(shape=(1000, 10))# 初始化並訓練PCA模型pca = PCA(n_components=2)pca.fit(data)# 取得PCA transformed_data = pca .transform( data)# 取得反向變換soriginal_data = pca.inverse_transform(transformed_data)# 分佈中的樣本X_sampled = pca.sample(n_samples=1000, key=None)
這仍然處於開發階段，效果很好，但預計會很粗糙，因此強烈鼓勵貢獻！
在不改變設計模式的情況下進行更改。
如有必要，為您的更改編寫測試。
使用pip3 install -e . 。
使用python3 -m unittest discover -s tests執行測試。
然後提交拉取請求。
貢獻可以透過多種形式進行：
編寫文檔。
修復錯誤。
實施文件。
編寫高覆蓋率測試。
優化現有程式碼。
實驗並向範例部分提交真實範例。
報告錯誤。
回應報告的問題。
加入 Discord 伺服器以了解更多資訊。
“NanoDL”這個名字代表奈米深度學習。模型的規模呈爆炸式增長，因此資源有限的專家和公司無法在無需高昂成本的情況下建立靈活的模型。隨著 Phi 模型的成功，長期目標是建立和訓練所有可用模型的奈米版本，同時確保它們在性能上與原始模型競爭，參數總數不超過 1B。經過訓練的權重將透過該庫提供。任何形式的贊助、資助都將有助於提供培訓資源。您可以透過 GitHub 進行贊助，也可以透過 [email protected] 聯繫。
引用這個儲存庫：
@software{nanodl2024github,
author = {Henry Ndubuaku},
title = {NanoDL: A Jax-based library for designing and training transformer models from scratch.},
url = {http://github.com/hmunachi/nanodl},
year = {2024},
} 
