flax 框架理解

哦豁灬

于 2024-05-21 23:06:39 发布

阅读量341

点赞数 10

分类专栏：深度学习文章标签：深度学习框架 flax 深度学习

本文链接：https://blog.csdn.net/qq_38342510/article/details/138543617

版权

深度学习专栏收录该内容

29 篇文章 1 订阅

订阅专栏

深度学习框架有很多，所有框架都要回答下面的几个基本问题：

如何定义网络？
如何初始化网络参数？
如何计算反向传播？
如何更新网络参数？
如何管理训练状态？

pytorch 作为越来越受欢迎的框架，以上几个问题的解决无疑是接近完美的，flax 相对于 pytorch，又是如何面临这几个问题的呢？

1 网络定义

flax采取就地定义，就地使用的方式，使用时再定义。

from flax import linen as nn
from flax.metrics import tensorboard
from flax.training import train_state
import jax
import jax.numpy as jnp
import ml_collections
import numpy as np
import optax
from sqlalchemy import false
import tensorflow_datasets as tfds

class CNN(nn.Module):
  @nn.compact
  def __call__(self, x,is_training:bool=True):
    x = nn.Conv(features=32, kernel_size=(3, 3))(x)
    x = nn.BatchNorm(use_running_average=not is_training)(x)
    x = nn.relu(x)
    x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
    x = nn.Conv(features=64, kernel_size=(3, 3))(x)
    x = nn.BatchNorm(use_running_average=not is_training)(x)
    x = nn.relu(x)
    x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
    x = x.reshape((x.shape[0], -1))  # flatten
    x = nn.Dense(features=256)(x)
    x = nn.relu(x)
    x = nn.Dense(features=10)(x)
    return x

2 初始化网络参数

使用网络的 init() 方法初始化网络参数，方法的参数需要输入数据的形状

  cnn = CNN()
  variables=cnn.init(rng, jnp.ones([1, 28, 28, 1]))
  params = variables['params']
  batch_stats=variables['batch_stats']

3 管理训练状态

TrainState.create创建训练状态，三个参数：前向传播函数，网络参数，优化器

tx = optax.sgd(0.01, 0.99)
state=train_state.TrainState.create(apply_fn=cnn.apply, params=params, tx=tx)

4 计算反向传播

先定义梯度计算函数grad_fn = jax.value_and_grad(loss_fn, has_aux=True)，实际就是grad函数和损失函数的复合函数；
调用grad_fn得到梯度，函数grad_fn 的参数与loss_fn一致，返回值就是grads ，结构与loss_fn的第一个参数params一致。
有两个函数可以计算梯度：jax.value_and_grad和jax.grad。

@jax.jit
def apply_model(state, images, labels,old_batch_stats):
  def loss_fn(params,old_batch_stats):
    logits,mutated_vars = state.apply_fn({'params': params,"batch_stats":old_batch_stats}, images,is_training=True, mutable=['batch_stats'])
    one_hot = jax.nn.one_hot(labels, 10)
    loss = jnp.mean(optax.softmax_cross_entropy(logits=logits, labels=one_hot))
    return loss, (logits,mutated_vars['batch_stats'])
    
  grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
  (loss, (logits,new_batch_stats)), grads = grad_fn(state.params,old_batch_stats)
  accuracy = jnp.mean(jnp.argmax(logits, -1) == labels)
  return grads, loss, accuracy,new_batch_stats

5 更新网络参数

训练状态state更新自己

state=state.apply_gradients(grads=grads)

完整代码：

from flax import linen as nn
from flax.metrics import tensorboard
from flax.training import train_state
import jax
import jax.numpy as jnp
import ml_collections
import numpy as np
import optax
from sqlalchemy import false
import tensorflow_datasets as tfds
import flax


class CNN(nn.Module):
  @nn.compact
  def __call__(self, x,is_training:bool=True):
    x = nn.Conv(features=32, kernel_size=(3, 3))(x)
    x = nn.BatchNorm(use_running_average=not is_training)(x)
    x = nn.relu(x)
    x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
    x = nn.Conv(features=64, kernel_size=(3, 3))(x)
    x = nn.BatchNorm(use_running_average=not is_training)(x)
    x = nn.relu(x)
    x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
    x = x.reshape((x.shape[0], -1))  # flatten
    x = nn.Dense(features=256)(x)
    x = nn.relu(x)
    x = nn.Dense(features=10)(x)
    return x



@jax.jit
def apply_model(state, images, labels,old_batch_stats):
  def loss_fn(params,old_batch_stats):
    logits,mutated_vars = state.apply_fn({'params': params,"batch_stats":old_batch_stats}, images,is_training=True, mutable=['batch_stats'])
    one_hot = jax.nn.one_hot(labels, 10)
    loss = jnp.mean(optax.softmax_cross_entropy(logits=logits, labels=one_hot))
    return loss, (logits,mutated_vars['batch_stats'])
    
  grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
  (loss, (logits,new_batch_stats)), grads = grad_fn(state.params,old_batch_stats)
  accuracy = jnp.mean(jnp.argmax(logits, -1) == labels)
  return grads, loss, accuracy,new_batch_stats


@jax.jit
def update_model(state, grads):
  return state.apply_gradients(grads=grads)


def train_epoch(state, train_ds, batch_size, rng,batch_stats):
  train_ds_size = len(train_ds['image'])
  steps_per_epoch = train_ds_size // batch_size

  perms = jax.random.permutation(rng, len(train_ds['image']))
  perms = perms[:steps_per_epoch * batch_size]  # skip incomplete batch
  perms = perms.reshape((steps_per_epoch, batch_size))

  epoch_loss = []
  epoch_accuracy = []
  for perm in perms:
    batch_images = train_ds['image'][perm, ...]
    batch_labels = train_ds['label'][perm, ...]
    grads, loss, accuracy ,batch_stats= apply_model(state, batch_images, batch_labels,batch_stats)
    
    state = update_model(state, grads)
    
    epoch_loss.append(loss)
    epoch_accuracy.append(accuracy)
  train_loss = np.mean(epoch_loss)
  train_accuracy = np.mean(epoch_accuracy)
  return state, train_loss, train_accuracy,batch_stats


def get_datasets():
  ds_builder = tfds.builder('fashion_mnist')
  ds_builder.download_and_prepare()
  train_ds = tfds.as_numpy(ds_builder.as_dataset(split='train', batch_size=-1))
  test_ds = tfds.as_numpy(ds_builder.as_dataset(split='test', batch_size=-1))
  train_ds['image'] = jnp.float32(train_ds['image']) / 255.
  test_ds['image'] = jnp.float32(test_ds['image']) / 255.
  return train_ds, test_ds


def create_train_state(rng):
  cnn = CNN()
  variables=cnn.init(rng, jnp.ones([1, 28, 28, 1]))
  params = variables['params']
  batch_stats=variables['batch_stats']
  tx = optax.sgd(0.01, 0.99)
  state=train_state.TrainState.create(apply_fn=cnn.apply, params=params, tx=tx)
  
  return state,batch_stats

@jax.jit
def predict(state, batch_stats,image_i):
  logits= state.apply_fn({'params': state.params,"batch_stats":batch_stats},image_i,is_training=False)
  return logits
def test(state, batch_stats,test_ds):
  images = test_ds['image']
  labels = test_ds['label']
  batchs=1000
  accuracy=0
  for i in range(0,len(images),batchs):
    image_i=images[i:i+batchs]
    label_i=labels[i:i+batchs]

    logits= predict(state, batch_stats,image_i)
    accuracy += jnp.sum(jnp.argmax(logits, -1) == label_i)
  
  return accuracy/len(images)

def train_and_evaluate() -> train_state.TrainState:
  train_ds, test_ds = get_datasets()
  rng = jax.random.PRNGKey(0)

  
  rng, init_rng = jax.random.split(rng)
  state,batch_stats = create_train_state(init_rng)
  
  for epoch in range(1, 100 + 1):
    rng, input_rng = jax.random.split(rng)
    state, train_loss, train_accuracy,batch_stats = train_epoch(state, train_ds,
                                                    100,
                                                    input_rng,batch_stats)
    
    
    print(test(state, batch_stats,test_ds),end=" ")
  
  return state

train_and_evaluate()