JAX 神經網路資料預處理與訓練最佳化

import jax.numpy as jnp

def preprocess_data(data, max_value):
    # 將整數值轉換為float32浮點數
    data = jnp.array(data, dtype=jnp.float32)
    
    # 將浮點數除以最大整數值
    data = data / max_value
    
    return data

import jax
from jax.experimental import jax2tf

# 定義批次大小
batch_size = 32

# 將資料分割為訓練集和測試集
train_data, test_data =...

# 建立批次生成器
train_batches = jax2tf.convert_batch(train_data, batch_size)
test_batches = jax2tf.convert_batch(test_data, batch_size)

import flax
from flax import linen as nn

class SimpleNeuralNetwork(nn.Module):
    @nn.compact
    def __call__(self, x):
        # 輸入層
        x = nn.Dense(512)(x)
        x = nn.relu(x)
        
        # 輸出層
        x = nn.Dense(10)(x)
        x = nn.softmax(x)
        
        return x

import jax
import jax.numpy as jnp

# 初始化模型引數
params = {
    'weights': jnp.random.normal(0, 1, (784, 512)),
    'biases': jnp.zeros((512,)),
    'output_weights': jnp.random.normal(0, 1, (512, 10)),
    'output_biases': jnp.zeros((10,))
}

# 定義前向傳遞函式
def predict(params, image):
    # 輸入層
    input_layer = image
    
    # 隱藏層
    hidden_layer = jnp.dot(input_layer, params['weights']) + params['biases']
    hidden_layer = jax.nn.relu(hidden_layer)
    
    # 輸出層
    output_layer = jnp.dot(hidden_layer, params['output_weights']) + params['output_biases']
    output_layer = jax.nn.softmax(output_layer)
    
    return output_layer

# 應用模型
image = jnp.random.normal(0, 1, (784,))
output = predict(params, image)
print(output)

import jax
import jax.numpy as jnp
from jax import random

# 定義層大小
LAYER_SIZES = [28*28, 512, 10]

# 定義引數縮放因子
PARAM_SCALE = 0.01

def init_network_params(sizes, key=random.PRNGKey(0), scale=1e-2):
    """初始化所有層的引數"""
    
    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,))
    
    # 初始化引數
    params = []
    for i in range(len(sizes) - 1):
        m, n = sizes[i], sizes[i + 1]
        w, b = random_layer_params(m, n, key, scale)
        params.append((w, b))
        key = random.fold_in(key, i)
    
    return params

# 初始化引數
params = init_network_params(LAYER_SIZES)

# 列印初始化引數
for i, (w, b) in enumerate(params):
    print(f"Layer {i+1} weights shape: {w.shape}, biases shape: {b.shape}")

import jax
import jax.numpy as jnp
from jax.experimental import jax2tf
from jax.experimental import stax

# 定義層大小
LAYER_SIZES = [784, 256, 10]

# 定義引數縮放因子
PARAM_SCALE = 0.1

# 生成隨機鍵
key = jax.random.PRNGKey(0)

# 將鍵分割為多個層所需的隨機鍵
keys = jax.random.split(key, len(LAYER_SIZES) - 1)

# 初始化神經網路引數
params = []
for m, n, k in zip(LAYER_SIZES[:-1], LAYER_SIZES[1:], keys):
    # 生成隨機權重和偏差
    w = jax.random.normal(k, (m, n)) * PARAM_SCALE
    b = jax.random.normal(k, (n,)) * PARAM_SCALE
    params.append((w, b))

import jax.nn as nn

def predict(params, x):
    # 逐層進行前向傳遞
    for w, b in params:
        x = nn.swish(jnp.dot(x, w) + b)
    return x

# 生成隨機輸入
x = jax.random.normal(key, (1, LAYER_SIZES[0]))

# 進行前向傳遞
output = predict(params, x)

import jax.numpy as jnp
from jax.nn import swish

def predict(params, image):
    """每個例子的預測函式"""
    activations = image
    for w, b in params[:-1]:
        outputs = jnp.dot(w, activations) + b
        activations = swish(outputs)
    final_w, final_b = params[-1]
    logits = jnp.dot(final_w, activations) + final_b
    return logits

from jax.nn import swish

import jax.numpy as jnp

# 初始化隨機影像
random_flattened_image = jnp.random.rand(784)

# 匯入神經網路模型
model =...

# 進行預測
prediction = model.predict(random_flattened_image)

print(prediction.shape)  # 輸出：(10,)

import jax
import jax.numpy as jnp
from jax import random

# 生成隨機批次影像
random_flattened_images = random.normal(random.PRNGKey(1), (32, 28*28*1))

try:
    preds = predict(params, random_flattened_images)
except TypeError as e:
    print(e)

# 使用vmap實作自動批次處理
batched_predict = jax.vmap(predict, in_axes=(None, 0))

# 測試批次預測
batched_preds = batched_predict(params, random_flattened_images)
print(batched_preds.shape)

from jax import vmap
batched_predict = vmap(predict, in_axes=(None, 0))

batched_preds = batched_predict(params, random_flattened_images)
print(batched_preds.shape)

(32, 10)

import jax
import jax.numpy as jnp

# 定義單個專案的函式
def predict(image):
    # 對影像進行預測
    return jnp.array([0.1, 0.2, 0.3, 0.4])  # 例如

# 使用 vmap() 建立批次化函式
batched_predict = jax.vmap(predict)

# 測試批次化函式
batch_images = jnp.array([[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]])
output = batched_predict(batch_images)
print(output)

# 定義損失函式
def loss(params, images, labels):
    # 對影像進行預測
    predictions = batched_predict(images)
    # 計算損失
    return jnp.mean((predictions - labels) ** 2)

# 初始化引數
params = jnp.array([0.1, 0.2, 0.3, 0.4])

# 定義最佳化器
def optimizer(params, gradient):
    # 更新引數
    return params - 0.01 * gradient

# 執行梯度下降迭代
for i in range(100):
    # 計算梯度
    gradient = jax.grad(loss)(params, batch_images, jnp.array([0, 1]))
    # 更新引數
    params = optimizer(params, gradient)
    # 印出損失
    print(loss(params, batch_images, jnp.array([0, 1])))

import numpy as np
import matplotlib.pyplot as plt

# 定義Loss函式
def loss(x):
    return x**2 + 10*np.sin(x)

# 生成x值
x = np.linspace(-10, 10, 400)

# 計算Loss值
y = loss(x)

# 繪製Loss曲線
plt.plot(x, y)
plt.xlabel('x')
plt.ylabel('Loss')
plt.title('Loss Curve')
plt.show()

import numpy as np

# 定義Loss函式和其梯度
def loss(x):
    return x**2 + 10*np.sin(x)

def gradient(x):
    return 2*x + 10*np.cos(x)

# 起始點
x = 5.0

# 學習率
lr = 0.1

# 更新步數
steps = 100

# 更新模型引數
for i in range(steps):
    # 計算梯度
    grad = gradient(x)
    
    # 更新模型引數
    x -= lr * grad
    
    # 輸出當前Loss值
    print(f'Step {i+1}, Loss: {loss(x)}')

import jax.numpy as jnp
from jax.nn import logsumexp

def loss(params, images, targets):
    """交叉熵損失函式"""
    logits = batched_predict(params, images)
    log_preds = logits - logsumexp(logits)
    return -jnp.mean(targets * log_preds)

from jax import grad

# 定義損失函式
def loss(params, images, targets):
   ...

# 取得梯度
grad_loss = grad(loss, argnums=0)

params -= learning_rate * grad_loss(params, images, targets)

import jax.numpy as jnp
from jax.nn import logsumexp
from jax import grad

def loss(params, images, targets):
    """交叉熵損失函式"""
    logits = batched_predict(params, images)
    log_preds = logits - logsumexp(logits)
    return -jnp.mean(targets * log_preds)

def gradient_descent(params, images, targets, learning_rate):
    """梯度下降法"""
    grad_loss = grad(loss, argnums=0)
    params -= learning_rate * grad_loss(params, images, targets)
    return params

# 初始化模型引數
params =...

# 載入資料
images =...
targets =...

# 設定學習率
learning_rate = 0.01

# 執行梯度下降法
for i in range(100):
    params = gradient_descent(params, images, targets, learning_rate)
    print(f"第 {i+1} 次迭代，損失值：{loss(params, images, targets)}")

from jax import grad

INIT_LR = 1.0  # 初始學習率
DECAY_RATE = 0.95  # 衰減率
DECAY_STEPS = 5  # 衰減步驟

def update(params, x, y, epoch_number):
    """
    更新模型引數。
    
    引數：
    - params: 模型引數（權重和偏差）
    - x: 輸入資料
    - y: 真實標籤
    - epoch_number: 當前epoch數
    
    傳回：
    - 更新後的模型引數
    """
    # 計算梯度
    grads = grad(loss)(params, x, y)
    
    # 計算當前epoch的學習率
    lr = INIT_LR * (DECAY_RATE ** (epoch_number / DECAY_STEPS))
    
    # 更新模型引數
    return [(w - lr * dw, b - lr * db) for (w, b), (dw, db) in zip(params, grads)]

from jax import value_and_grad

def update(params, x, y, epoch_number):
    """
    更新模型引數，並計算損失值和梯度。
    
    引數：
    - params: 模型引數（權重和偏差）
    - x: 輸入資料
    - y: 真實標籤
    - epoch_number: 當前epoch數
    
    傳回：
    - 更新後的模型引數、損失值和梯度
    """
    # 計算損失值和梯度
    loss_value, grads = value_and_grad(loss)(params, x, y)
    
    # 計算當前epoch的學習率
    lr = INIT_LR * (DECAY_RATE ** (epoch_number / DECAY_STEPS))
    
    # 更新模型引數
    updated_params = [(w - lr * dw, b - lr * db) for (w, b), (dw, db) in zip(params, grads)]
    
    return updated_params, loss_value

import numpy as np

def value_and_grad(loss):
    # 自動微分函式，計算梯度
    def grad_func(params, x, y):
        # 計算梯度
        loss_value = loss(params, x, y)
        grads = np.gradient(loss_value)
        return loss_value, grads
    return grad_func

# 初始化引數和學習率
params = (np.random.rand(), np.random.rand())
x = np.random.rand()
y = np.random.rand()
INIT_LR = 0.1
DECAY_RATE = 0.9
DECAY_STEPS = 10

# 計算梯度和更新引數
loss_value, grads = value_and_grad(lambda params, x, y: np.mean((params[0]*x + params[1] - y)**2))(params, x, y)
lr = INIT_LR * DECAY_RATE ** (1 / DECAY_STEPS)
params_updated = (params[0] - lr * grads[0], params[1] - lr * grads[1])

print("更新後的引數：", params_updated)

import jax.numpy as jnp
from jax.nn import one_hot

# 定義訓練迴圈的引數
num_epochs = 25  # 訓練迴圈的次數

# 定義批次準確率計算函式
def batch_accuracy(params, images, targets):
    """
    計算批次的準確率。
    
    Args:
    params: 模型引數
    images: 輸入影像
    targets: 目標標籤
    
    Returns:
    準確率
    """
    images = jnp.reshape(images, (len(images), NUM_PIXELS))  # 重塑影像為 (batch_size, NUM_PIXELS)
    predicted_class = jnp.argmax(batched_predict(params, images), axis=1)  # 預測類別
    return jnp.mean(predicted_class == targets)  # 準確率

# 定義整體準確率計算函式
def accuracy(params, data):
    """
    計算整體準確率。
    
    Args:
    params: 模型引數
    data: 資料集
    
    Returns:
    準確率
    """
    accs = []  # 準確率列表
    for images, targets in data:
        accs.append(batch_accuracy(params, images, targets))  # 計算批次準確率
    return jnp.mean(jnp.array(accs))  # 整體準確率

import time  # 引入時間模組

# 訓練迴圈
for epoch in range(num_epochs):
    start_time = time.time()  # 記錄開始時間
    
    losses = []  # 損失列表
    
    for x, y in train_data:
        x = jnp.reshape(x, (len(x), NUM_PIXELS))  # 重塑輸入影像
        
        y = one_hot(y, NUM_LABELS)  # 對目標標籤進行 one-hot 編碼
        
        params, loss_value = update(params, x, y, epoch)  # 更新模型引數
        
        losses.append(loss_value)  # 記錄損失值
    
    epoch_time = time.time() - start_time  # 記錄本次迭代的時間
    
    # 可以在這裡新增列印或儲存相關資訊的程式碼

# 計算訓練集上的準確率
train_acc = accuracy(params, train_data)

# 計算測試集上的準確率
test_acc = accuracy(params, test_data)

# 記錄開始時間
start_time = time.time()

# 執行評估
eval_time = time.time() - start_time

# 列印評估時間
print("Eval in {:0.2f} sec".format(eval_time))

# 定義損失函式和梯度計算
def calculate_loss_and_gradient(params, data):
    # 計算損失函式
    loss = calculate_loss(params, data)
    
    # 計算梯度
    gradient = calculate_gradient(params, data)
    
    return loss, gradient

# 更新模型引數
def update_params(params, gradient):
    # 更新引數
    updated_params = params - learning_rate * gradient
    
    return updated_params

# 定義one-hot編碼函式
def one_hot_encoding(label):
    # 建立one-hot編碼向量
    one_hot_vector = np.zeros((num_classes,))
    one_hot_vector[label] = 1
    
    return one_hot_vector

# 定義訓練模型函式
def train_model(params, train_data, num_epochs):
    for epoch in range(num_epochs):
        # 前向傳播
        outputs = forward_propagate(params, train_data)
        
        # 計算損失函式和梯度
        loss, gradient = calculate_loss_and_gradient(params, train_data)
        
        # 更新模型引數
        params = update_params(params, gradient)
        
        # 列印訓練過程資訊
        print("Epoch {} in {:0.2f} sec".format(epoch, epoch_time))
        print("Training set loss {}".format(jnp.mean(jnp.array(losses))))
        print("Training set accuracy {}".format(train_acc))
        print("Test set accuracy {}".format(test_acc))