具狀態的計算

具狀態的計算#

JAX 轉換，例如 jit()、vmap()、grad()，要求它們封裝的函式必須是純函式：也就是說，輸出僅完全依賴於輸入，並且沒有副作用（例如更新全域狀態）的函式。您可以在 JAX 尖銳之處：純函式中找到相關討論。

在機器學習的背景下，這種限制可能會帶來一些挑戰，因為狀態可能以多種形式存在。例如：

模型參數，
最佳化器狀態，以及
具狀態的層，例如 BatchNorm。

本節提供關於如何在 JAX 程式中正確處理狀態的一些建議。

簡單範例：計數器#

讓我們從查看一個簡單的具狀態程式開始：計數器。

import jax
import jax.numpy as jnp

class Counter:
  """A simple counter."""

  def __init__(self):
    self.n = 0

  def count(self) -> int:
    """Increments the counter and returns the new value."""
    self.n += 1
    return self.n

  def reset(self):
    """Resets the counter to zero."""
    self.n = 0


counter = Counter()

for _ in range(3):
  print(counter.count())

1
2
3

計數器的 n 屬性在連續呼叫 count 之間維護計數器的狀態。它會因為呼叫 count 的副作用而被修改。

假設我們想要快速計數，因此我們對 count 方法進行 JIT 編譯。（在本範例中，這實際上並不能提高速度，原因有很多，但請將此視為 JIT 編譯模型參數更新的玩具模型，其中 jit() 會產生巨大的差異）。

counter.reset()
fast_count = jax.jit(counter.count)

for _ in range(3):
  print(fast_count())

1
1
1

糟糕！我們的計數器無法運作。這是因為這行程式碼

self.n += 1

在 count 中涉及副作用：它會就地修改輸入計數器，因此 jit 不支援此函式。此類副作用僅在首次追蹤函式時執行一次，後續呼叫不會重複副作用。那麼，我們該如何修正呢？

解決方案：顯式狀態#

我們計數器的部分問題是傳回值不依賴於引數，這表示常數「烘焙」到編譯後的輸出中。但不應該是常數 – 它應該依賴於狀態。那麼，我們為什麼不將狀態變成引數呢？

CounterState = int

class CounterV2:

  def count(self, n: CounterState) -> tuple[int, CounterState]:
    # You could just return n+1, but here we separate its role as 
    # the output and as the counter state for didactic purposes.
    return n+1, n+1

  def reset(self) -> CounterState:
    return 0

counter = CounterV2()
state = counter.reset()

for _ in range(3):
  value, state = counter.count(state)
  print(value)

1
2
3

在這個新版本的 Counter 中，我們將 n 移至 count 的引數，並新增了另一個傳回值，代表新的、已更新的狀態。若要使用此計數器，我們現在需要顯式追蹤狀態。但作為回報，我們現在可以安全地 jax.jit 這個計數器

state = counter.reset()
fast_count = jax.jit(counter.count)

for _ in range(3):
  value, state = fast_count(state)
  print(value)

1
2
3

通用策略#

我們可以將相同的流程應用於任何具狀態的方法，以將其轉換為無狀態的方法。我們採用了以下形式的類別：

class StatefulClass

  state: State

  def stateful_method(*args, **kwargs) -> Output:

並將其轉換為以下形式的類別：

class StatelessClass

  def stateless_method(state: State, *args, **kwargs) -> (Output, State):

這是一種常見的函式式程式設計模式，而且基本上是所有 JAX 程式中處理狀態的方式。

請注意，一旦我們以這種方式重寫它，對類別的需求就變得不太明確了。由於類別不再執行任何工作，因此我們可以僅保留 stateless_method。這是因為，就像我們剛才應用的策略一樣，物件導向程式設計 (OOP) 是一種幫助程式設計師理解程式狀態的方式。

在我們的例子中，CounterV2 類別只不過是一個命名空間，將所有使用 CounterState 的函式集中在一個位置。讀者的練習：您認為將其保留為類別是否有意義？

順帶一提，您已經在 JAX 虛擬隨機性 API jax.random 中看過此策略的範例，如虛擬隨機數節所示。與使用隱式更新的具狀態類別來管理隨機狀態的 Numpy 不同，JAX 要求程式設計師直接使用隨機產生器狀態 – PRNG 金鑰。

簡單的範例：線性迴歸#

讓我們將此策略應用於一個簡單的機器學習模型：透過梯度下降的線性迴歸。

在這裡，我們僅處理一種狀態：模型參數。但通常，您會看到許多種類的狀態被串連到 JAX 函式中和從 JAX 函式中串連出來，例如最佳化器狀態、batchnorm 的層統計資料等。

要仔細查看的函式是 update。

from typing import NamedTuple

class Params(NamedTuple):
  weight: jnp.ndarray
  bias: jnp.ndarray


def init(rng) -> Params:
  """Returns the initial model params."""
  weights_key, bias_key = jax.random.split(rng)
  weight = jax.random.normal(weights_key, ())
  bias = jax.random.normal(bias_key, ())
  return Params(weight, bias)


def loss(params: Params, x: jnp.ndarray, y: jnp.ndarray) -> jnp.ndarray:
  """Computes the least squares error of the model's predictions on x against y."""
  pred = params.weight * x + params.bias
  return jnp.mean((pred - y) ** 2)


LEARNING_RATE = 0.005

@jax.jit
def update(params: Params, x: jnp.ndarray, y: jnp.ndarray) -> Params:
  """Performs one SGD update step on params using the given data."""
  grad = jax.grad(loss)(params, x, y)

  # If we were using Adam or another stateful optimizer,
  # we would also do something like
  #
  #   updates, new_optimizer_state = optimizer(grad, optimizer_state)
  # 
  # and then use `updates` instead of `grad` to actually update the params.
  # (And we'd include `new_optimizer_state` in the output, naturally.)

  new_params = jax.tree_map(
      lambda param, g: param - g * LEARNING_RATE, params, grad)

  return new_params

請注意，我們手動將參數傳入和傳出更新函式。

import matplotlib.pyplot as plt

rng = jax.random.key(42)

# Generate true data from y = w*x + b + noise
true_w, true_b = 2, -1
x_rng, noise_rng = jax.random.split(rng)
xs = jax.random.normal(x_rng, (128, 1))
noise = jax.random.normal(noise_rng, (128, 1)) * 0.5
ys = xs * true_w + true_b + noise

# Fit regression
params = init(rng)
for _ in range(1000):
  params = update(params, xs, ys)

plt.scatter(xs, ys)
plt.plot(xs, params.weight * xs + params.bias, c='red', label='Model Prediction')
plt.legend();

/tmp/ipykernel_2808/721844192.py:37: DeprecationWarning: jax.tree_map is deprecated: use jax.tree.map (jax v0.4.25 or newer) or jax.tree_util.tree_map (any JAX version).
  new_params = jax.tree_map(

_images/e84ae8938e7347c77e584740263882f4331422be79c00e90c4f90970603133b5.png

更進一步#

上述策略是任何 JAX 程式在使用 jit、vmap、grad 等轉換時必須處理狀態的方式。

如果您處理的是兩個參數，則手動處理參數似乎還可以，但如果它是具有數十層的神經網路呢？您可能已經開始擔心以下兩件事：

我們是否應該手動初始化所有參數，基本上重複我們已經在前向傳遞定義中寫入的內容？
我們是否應該手動串連所有這些東西？

細節可能很棘手，但有一些程式庫範例可以為您處理這些問題。請參閱 JAX 神經網路程式庫以取得一些範例。