jax和stax学习_44实用知识库分享.pdfVIP

jax和stax学习_44 目录 jax和stax学习_44 Jax:一个比numpy、pytorch、cupy还要理想的工具 jax：自定义损失函数 8 Jax专家太有爱了，Nice！ stax：Jax的示例神经网络库 Jax中的卷积 19 Jax：用stax写一个卷积神经网络 JAX101: 认识Pytrees数据结构 flax中BatchNormal的使用 jax:循环结构fori_loop和while_loop jax：自定义激活函数 43 jax：使用 PyTorch 加载数据 jax：optax优化库 jax:用SCAN构建一个简单的LSTM网络 jax:flax模型保存与加载 jax：CTC loss 实作与优化 flax：一个很好用的jax框架 jax：pytree基本数据类型 flax：BatchNorm和Dropout的使用 Jax:一个比numpy、pytorch、cupy还要理想的工具 Jax:一个比numpy、pytorch、cupy还要理想的工具 一、理想中的工具是足够简单，又足够强大。 numpy是，可惜numpy不能自动微分，也不能调用GPU； cupy兼容numpy，也能调用GPU，但还是不能自动微分； pytorch强大而稳定可靠，但与numpy不兼容，上来就要符合他的编程模型和框架，还不足够简 单； Jax来了，他与numpy兼容，还能调用GPU、TPU，并行运算、还能自动微分，太完美了！ Jax官网介绍： 二、思维模型：复合函数 Jax三个关键的基础设施：自动微分、GPU加速、向量化运算，分别对应三个函数grad、jit、 vmap. 还是以mnist为例 1、我们随便定义一个损失函数loss： def loss (param,batch): inputs_list, targets_list=batch final_outputs=predict (param,inputs_list) t= -targets_list *jnp.log (final_outputs) loss=jnp.sum(t) return loss 2、复合上面的三个函数，就可以得到新的函数 backward=jit (grad (loss)) 3、调用backward： 学过复合函数的同学都知道，复合函数调用参数跟最里层函数一样，返回值跟最外层函数一样： 那么返回值是什么呢？ jit只是GPU加速，不改变返回值； grad的返回值是param的反向传播梯度，结构与param（即loss函数的第一个参数）一样。 第 1 页 /共 85 页 Jax:一个比numpy、pytorch、cupy还要理想的工具 这里param的结构其实并不是jnp.array的shape，而是由list、tuple、dict相互嵌套构成的结 构，grad会遍历这个结构。除了list、tuple、dict，也是可以的，但不能是python的普通类。 grad不会遍历jnp.array。而是把array整体当成一个值来看，就是pytree的叶子节点，这个值当 然就有shape了，得到的梯度值与它的shape是相同的。 不过这样一搞，loss函数就不能是类的函数了，类函数第一个参数是self（哭），所以要能用 grad功能的，只能是纯函数。 所谓纯函数，就是不能像类一样带自身的状态参数，所有值都只能由函数参数输入，如果要用全 局变量，也只能是静态的全局变量。 4、定义参数 这里是numpy和Jax混合编程了，很融洽，感觉没有特别需要遵守的规矩 wih = jnp.array (numpy.random.normal (0.0, pow (inodes, -0.5), (784,400))) who = jnp.array (numpy.random.normal (0.0, pow (hnodes