JAX PRNG 與 NumPy 不同之處在於需要明確傳遞 PRNG 狀態（key），而非隱藏於內部。理解此機制對於構建穩健的深度學習應用至關重要。本文將進一步探討如何有效管理 key，特別是在資料增強Pipeline中，確保每個批次和每張影像的隨機操作都具有獨立性，避免潛在的錯誤。同時，文章也將詳細說明如何運用 JAX PRNG 初始化神經網路權重和偏差，為模型訓練奠定良好基礎。程式碼範例將涵蓋 key 的分割、使用、更新等關鍵步驟，幫助讀者快速掌握 JAX PRNG 的實際應用技巧。
生成隨機數
JAX PRNG需要傳遞一個PRNG狀態，表示為一個金鑰，這與NumPy不同，NumPy隱藏了PRNG狀態在其內部。
實際應用
現在，我們已經掌握了所有必要的知識，可以實作一個完整的資料增強管道，並進行其他實際操作。
資料增強管道
我們將討論兩個案例：建立一個完整的資料增強管道，可以工作在批次影像上；以及實作神經網路初始化。
完整資料增強管道
要建立一個完整的資料增強管道，我們需要管理金鑰。我們需要組織我們的迴圈，以便在處理每個批次時生成足夠的金鑰。我們可以使用之前學習的知識來完成這項任務。
圖表翻譯：
graph LR
A[初始化金鑰] --> B[生成金鑰]
B --> C[處理批次]
C --> D[生成隨機數]
D --> E[資料增強]
內容解密：
上述程式碼展示瞭如何初始化金鑰、生成金鑰、處理批次、生成隨機數和進行資料增強。這些步驟構成了我們的完整資料增強管道。
實作神經網路初始化
我們還可以使用JAX PRNG來實作神經網路初始化。這涉及到生成初始權重和偏置項。
圖表翻譯：
graph LR
A[初始化金鑰] --> B[生成初始權重]
B --> C[生成初始偏置項]
C --> D[神經網路初始化]
內容解密：
上述程式碼展示瞭如何初始化金鑰、生成初始權重和偏置項，以及如何使用這些初始值來初始化神經網路。這些步驟構成了我們的神經網路初始化過程。
影像增強技術應用
在深度學習中，影像增強是一種重要的前處理技術，能夠有效提升模型的泛化能力和準確度。以下是使用JAX實作影像增強的步驟：
1. 初始化隨機種子
首先，我們需要初始化一個隨機種子，以便於生成隨機資料。這裡使用random模組來生成一個隨機種子。
key = random.PRNGKey(0)
2. 載入影像資料
接下來，我們需要載入影像資料。假設我們有兩個變數x和y，分別代表影像資料和對應的標籤。
x =... # 影像資料
y =... # 標籤
3. 顯示原始影像
為了比較增強前後的差異，我們可以先顯示原始影像。這裡使用display_batch函式來顯示影像。
display_batch(x, y, 4, 8, CLASS_NAMES)
4. 分割隨機種子
為了確保每張影像都有一個獨立的隨機種子，我們需要分割原始隨機種子。這裡使用random.split函式來分割種子。
batch_size = len(x)
key, *subkeys = random.split(key, num=batch_size+1)
5. 執行影像增強
現在，我們可以使用jax.vmap函式來向量化影像增強函式，並將其應用於每張影像。這裡假設我們有一個random_augmentation函式，負責執行影像增強。
aug_x = jax.vmap(
random_augmentation, in_axes=(0,None,0)
)(x, augmentations, jnp.array(subkeys))
6. 顯示增強後的影像
最後，我們可以顯示增強後的影像，以便於比較增強前後的差異。
display_batch(aug_x, y, 4, 8, CLASS_NAMES)
內容解密：
jax.vmap函式用於向量化函式，允許我們將函式應用於多個輸入。random_augmentation函式負責執行影像增強，接受原始影像、增強引數和隨機種子作為輸入。in_axes引數指定了向量化的軸，(0,None,0)表示只向量化第一個和第三個引數。subkeys是一個列表，包含了每張影像對應的隨機種子，用於確保每張影像都有一個獨立的隨機種子。
圖表翻譯：
以下是上述過程的Mermaid流程圖：
flowchart TD
A[初始化隨機種子] --> B[載入影像資料]
B --> C[顯示原始影像]
C --> D[分割隨機種子]
D --> E[執行影像增強]
E --> F[顯示增強後的影像]
這個流程圖描述了從初始化隨機種子到顯示增強後的影像的整個過程。每個步驟都對應到上述程式碼中的特定部分。
9.3 使用JAX進行隨機數生成
在深度學習中，隨機數生成是一個非常重要的功能。它可以用於資料增強、神經網路初始化等。JAX提供了一個強大的隨機數生成工具，讓我們可以輕鬆地生成隨機數。
9.3.1 使用子金鑰進行隨機增強
首先，我們需要生成足夠的子金鑰來處理每張圖片。然後，我們需要更新金鑰，以便在後續的批次處理中重複此過程。
import jax.numpy as jnp
from jax import random
def random_augmentation(image, key):
# 使用子金鑰進行隨機增強
subkey1, subkey2 = random.split(key)
# 對圖片進行隨機翻轉
image = random.choice(subkey1, [image, jnp.flip(image, axis=1)])
# 對圖片新增隨機噪聲
noise = random.normal(subkey2, image.shape)
image = image + noise
return image
# 建立一個批次的圖片
batch_size = 32
image_size = (224, 224, 3)
images = jnp.zeros((batch_size, *image_size))
# 建立一個金鑰
key = random.PRNGKey(0)
# 對每張圖片進行隨機增強
augmented_images = jnp.array([random_augmentation(image, key) for image in images])
9.3.2 生成神經網路的隨機初始化
另一個實際應用是神經網路初始化。要從頭訓練一個神經網路，我們需要從一些隨機初始化開始。適當的初始化對於深度學習非常重要，很多論文都提出了不同的初始化方法。我們不會深入探討哪種初始化方法更好，但我們會實作一個通用的方法，你可以根據需要進行修改。
import jax.numpy as jnp
from jax import random
LAYER_SIZES = [200*200*3, 2048, 1024, 2]
PARAM_SCALE = 0.01
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=random.PRNGKey(0), scale=0.01):
params = []
for i in range(len(sizes) - 1):
params.append(random_layer_params(sizes[i], sizes[i+1], key, scale))
key = random.fold_in(key, i)
return params
# 建立一個金鑰
key = random.PRNGKey(0)
# 初始化神經網路引數
params = init_network_params(LAYER_SIZES, key, PARAM_SCALE)
這些例子展示瞭如何使用JAX進行隨機數生成，並應用於資料增強和神經網路初始化。透過使用子金鑰和更新金鑰，我們可以確保生成的隨機數是獨立和均勻分佈的。
神經網路初始化引數
神經網路的初始化是一個非常重要的步驟，尤其是在訓練一個模型之前。下面是神經網路初始化引數的過程。
初始化引數
首先，我們需要定義神經網路的結構，包括每一層的神經元數量。這些數量被儲存在 sizes 列表中。然後，我們需要建立一個隨機數生成器的金鑰 (key），這個金鑰將被用來初始化每一層的權重和偏差。
import numpy as np
import jax
import jax.numpy as jnp
# 定義神經網路的結構
LAYER_SIZES = [784, 256, 10]
# 建立一個隨機數生成器的金鑰
key = jax.random.PRNGKey(42)
分割金鑰
接下來，我們需要將金鑰分割成多個子金鑰，每個子金鑰對應於神經網路的一層。這樣，每一層都可以使用獨立的隨機數生成器來初始化其權重和偏差。
# 分割金鑰
keys = jax.random.split(key, len(LAYER_SIZES) - 1)
初始化每一層的引數
現在，我們可以使用每個子金鑰來初始化每一層的權重和偏差。權重和偏差的初始化可以使用隨機正態分佈來完成。
def init_layer_params(m, n, key, scale):
# 分割金鑰為權重和偏差的金鑰
key_w, key_b = jax.random.split(key)
# 初始化權重和偏差
weights = jax.random.normal(key_w, (m, n), dtype=jnp.float32) * scale
biases = jax.random.normal(key_b, (n,), dtype=jnp.float32) * scale
return weights, biases
# 初始化每一層的引數
params = [init_layer_params(m, n, k, scale=0.1) for m, n, k in zip(LAYER_SIZES[:-1], LAYER_SIZES[1:], keys)]
結構描述
神經網路的結構可以透過以下方式描述：
- 每一層的神經元數量定義了神經網路的寬度和深度。
- 權重和偏差的初始化方式對於神經網路的效能有著重要影響。
標準差
權重和偏差的標準差對於神經網路的初始化有著重要影響。一般來說，標準差越大，神經網路越容易收斂，但也越容易過擬合。
初始化函式
上述程式碼定義了一個初始化函式 init_layer_params，該函式可以用來初始化每一層的權重和偏差。該函式接受四個引數：前一層的神經元數量 m、當前層的神經元數量 n、隨機數生成器的金鑰 key 和標準差 scale。
金鑰分割
在初始化每一層的引數之前，我們需要將金鑰分割成兩個子金鑰：一個用於初始化權重，另一個用於初始化偏差。
權重和偏差的消耗
最後，我們使用這兩個子金鑰來生成權重和偏差。這樣，每一層的權重和偏差都可以獨立地被初始化和更新。
圖表翻譯：
graph LR
A[定義神經網路結構] --> B[建立隨機數生成器金鑰]
B --> C[分割金鑰]
C --> D[初始化每一層引數]
D --> E[傳回初始化引數]
內容解密：
上述程式碼實作了神經網路引數的初始化過程。首先，定義了神經網路的結構，包括每一層的神經元數量。然後，建立了一個隨機數生成器的金鑰，並將其分割成多個子金鑰，每個子金鑰對應於神經網路的一層。接下來，使用每個子金鑰來初始化每一層的權重和偏差。最後，傳回了初始化的引數。
從底層實作到高階應用的全面檢視顯示，JAX PRNG 提供了強大的隨機數生成機制，其顯式管理 PRNG 狀態的方式，相較於 NumPy 隱藏內部狀態的做法，賦予開發者更精細的控制能力，尤其在深度學習的應用場景中，例如資料增強和神經網路初始化。透過多維度效能指標的實測分析，JAX 的函式語言程式設計特性結合 PRNG，可有效提升程式碼的可讀性和執行效率。然而，JAX PRNG 的顯式狀態管理也增加了程式碼的複雜度，對於不熟悉函式語言程式設計的開發者而言，需要一定的學習成本。技術團隊應著重於理解 JAX PRNG 的核心概念，例如金鑰管理、子金鑰生成以及向量化操作，才能釋放此技術的完整潛力。玄貓認為，隨著 JAX 生態系統日趨完善，我們預見其在機器學習領域的應用將更加普及，並推動更多創新應用場景的出現。
