销售管理软件永久免费,seo翻译,软件开发人,重庆百度总代理目前来看Google 是唯一一家在 AI 价值链上实现端到端垂直整合的公司。从基础模型 (Gemini)、应用层 (ImageFX, Search with Gemini, NotebookLM)#xff0c;到云架构 (Google Cloud, Vertex AI) 以及硬件 (TPUs)#xff0c;几乎全都有所布局。长期以来Google 一直在通过提升自…目前来看Google 是唯一一家在 AI 价值链上实现端到端垂直整合的公司。从基础模型 (Gemini)、应用层 (ImageFX, Search with Gemini, NotebookLM)到云架构 (Google Cloud, Vertex AI) 以及硬件 (TPUs)几乎全都有所布局。长期以来Google 一直在通过提升自身能力来减少对 NVIDIA GPU 的依赖。这种技术积累逐渐演变成了现在的 JAX AI 栈。更有意思的是这套技术栈现在不仅 Google 自己用Anthropic、xAI 甚至 Apple 这些头部 LLM 提供商也都在用所以我们就很有必要这就很有必要深入聊聊这套技术栈了。什么是 JAX AI 栈简单来说JAX AI 栈是一套面向超大规模机器学习的端到端开源平台。核心组件主要由以下四个部分构成1、JAXGoogle 和 NVIDIA 联合开发的 Python 高性能数值计算库。接口设计极其类似 NumPy但区别在于它能自动、高效地在 CPU、GPU 或 TPU 上运行无论是本地还是分布式环境。底层的技术在于XLA (Accelerated Linear Algebra)编译器它能把 JAX 代码转译成针对不同硬件深度优化的机器码。对比之下NumPy 的操作默认只能在 CPU 上跑效率天差地别。2、Flax基于 JAX 的神经网络训练库。Flax 的核心现在是 NNX (NeuralNetworks for JAX)。这是一个简化版的 API让创建、调试和分析 JAX 神经网络变得更直观。之前有个 Flax Linen是那种无状态、函数式风格的 API。而 NNX 作为继任者引入了面向对象和有状态的特性对于习惯了 PyTorch 的开发者来说构建和调试 JAX 模型会顺手很多。3、OptaxJAX 生态里的梯度处理和优化库。它的优势在于灵活性几行代码就能把标准优化器和复杂的技巧比如梯度裁剪、梯度累积链式组合起来。4. Orbax专门处理 Checkpoint 的库用于保存和恢复大规模训练任务。支持异步分布式检查点这在大模型训练里至关重要——万一硬件挂了能从断点恢复不至于让昂贵的算力打了水漂。下面这张图展示了完整栈的架构除了上面这四个核心还有很多其他组件建议细看。实战用 JAX 训练神经网络JAX 之所以在 GPU 和 TPU 上能跑赢 PyTorch主要归功于即时 (JIT) 优化和 XLA 的后端编译效率。我们直接上手用 JAX 撸一个简单的神经网络搞个手写数字识别看看这套栈在实际工作流里到底怎么用。1、环境配置JAX AI 栈现在整合成了一个 metapackage安装很简单。然后我们还需要sklearn加载数据和matplotlib画图。!uv pip install jax-ai-stack sklearn matplotlib2、加载数据直接用 sklearn 加载 UCI ML 手写数字数据集。fromsklearn.datasetsimportload_digits # Load dataset digitsload_digits()数据是8 x 8的像素化手写数字图像0 到 9及其对应的标签。print(fNumber of samples × features: {digits.data.shape}) print(fNumber of labels: {digits.target.shape}) Number of samples × features: (1797, 64) Number of labels: (1797,) 3、 数据可视化先看看数据长什么样挑 100 张图画出来。import matplotlib.pyplot as plt fig, axes plt.subplots(10, 10, figsize(6, 6), subplot_kw{xticks:[], yticks:[]}, gridspec_kwdict(hspace0.1, wspace0.1)) for i, ax in enumerate(axes.flat): ax.imshow(digits.images[i], cmapbinary, interpolationgaussian) ax.text(0.05, 0.05, str(digits.target[i]), transformax.transAxes, colorgreen)4、 数据集切分常规操作把数据切成训练集和测试集。from sklearn.model_selection import train_test_split # Create dataset splits splits train_test_split(digits.images, digits.target, random_state0)5、转为 JAX 数组这一步很关键输入到模型之前需要用 JAX Numpy 把数据转成 JAX 数组格式。import jax.numpy as jnp # Convert splits to JAX arrays images_train, images_test, label_train, label_test map(jnp.asarray, splits)看一眼数据维度print(fTraining images shape: {images_train.shape}) print(fTraining labels shape: {label_train.shape}) print(fTest images shape: {images_test.shape}) print(fTest labels shape: {label_test.shape}) Training images shape: (1347, 8, 8) Training labels shape: (1347,) Test images shape: (450, 8, 8) Test labels shape: (450,) 6、用 Flax 构建网络用 Flax NNX 搭建一个带 SELU 激活函数的简单前馈网络。习惯写 PyTorch 的朋友会发现这语法看着非常眼熟。from flax import nnx class DigitClassifier(nnx.Module): def __init__(self, n_features, n_hidden, n_targets, rngs): self.n_features n_features self.layer_1 nnx.Linear(n_features, n_hidden, rngs rngs) self.layer_2 nnx.Linear(n_hidden, n_hidden, rngs rngs) self.layer_3 nnx.Linear(n_hidden, n_targets, rngs rngs) def __call__(self, x): x x.reshape(x.shape[0], self.n_features) [#Flatten](#Flatten) images x nnx.selu(self.layer_1(x)) x nnx.selu(self.layer_2(x)) x self.layer_3(x) return x7、实例化模型JAX 处理随机数的方式比较特别。这里用nnx.Rngs(0)初始化一个种子为 0 的随机数生成器 (RNG) 对象。这个对象负责管理网络操作里的所有随机性比如参数初始化和 Dropout。注意这和 PyTorch 直接设全局种子torch.manual_seed(seed)的逻辑不一样。# Initialize random number generator rngs nnx.Rngs(0) # Create instance of the classifier model DigitClassifier(n_features64, n_hidden128, n_targets10, rngs rngs)8、定义优化器与训练步骤用 Optax 定义优化器和损失函数。import jax import optax # SGD optimizer with learning rate 0.05 optimizer nnx.ModelAndOptimizer( model, optax.sgd(learning_rate0.05)) # Loss function def loss_fn(model, data, labels): # Forward pass logits model(data) # Compute mean cross-entropy loss loss optax.softmax_cross_entropy_with_integer_labels( logitslogits, labelslabels).mean() return loss, logits # Single training step with automatic differentiation and optimization nnx.jit # JIT compile for faster execution def training_step(model, optimizer, data, labels): loss_gradient nnx.grad(loss_fn, has_auxTrue) # has_auxTrue allows returning auxiliary outputs (logits) grads, logits loss_gradient(model, data, labels) # Forward backward pass optimizer.update(grads) # Update model parameters using computed gradients代码里用到了两个核心变换这是 JAX 高效的秘诀jax.jit即时编译把训练函数扔给 XLA 编译器重复执行速度极快。jax.grad利用自动微分计算梯度。Flax NNX 把它俩封装成了装饰器nnx.jit和nnx.grad用起来更方便。9、训练循环跑 500 epoch每 100 轮显示 Loss。num_epochs500 print_every100 forepochinrange(num_epochs1): # Training step training_step(model, optimizer, images_train, label_train) # Evaluate and print metrics periodically ifepoch%print_every0: train_loss, _loss_fn(model, images_train, label_train) test_loss, _loss_fn(model, images_test, label_test) print(fEpoch {epoch:3d} | Train Loss: {train_loss:.4f} | Test Loss: {test_loss:.4f}) Epoch 0 | Train Loss: 0.0044 | Test Loss: 0.1063 Epoch 100 | Train Loss: 0.0035 | Test Loss: 0.1057 Epoch 200 | Train Loss: 0.0029 | Test Loss: 0.1054 Epoch 300 | Train Loss: 0.0024 | Test Loss: 0.1052 Epoch 400 | Train Loss: 0.0021 | Test Loss: 0.1051 Epoch 500 | Train Loss: 0.0019 | Test Loss: 0.1050 10. 效果评估最后看看在测试集上的表现。# Evaluate model accuracy on test set logits model(images_test) predictions logits.argmax(axis1) correct jnp.sum(predictions label_test) total len(label_test) accuracy correct / total print(fTest Accuracy: {correct}/{total} correct ({accuracy:.2%})) # Test Accuracy: 437/450 correct (97.11%)97% 的准确率对于这么简单的网络来说相当不错了。最后把预测结果可视化一下绿色是对的红色是错的。fig, axes plt.subplots(10, 10, figsize(6, 6), subplot_kw{xticks:[], yticks:[]}, gridspec_kwdict(hspace0.1, wspace0.1)) for i, ax in enumerate(axes.flat): ax.imshow(images_test[i], cmapbinary, interpolationgaussian) color green if label_pred[i] label_test[i] else red ax.text(0.05, 0.05, str(label_pred[i]), transformax.transAxes, colorcolor)到这里你就已经在 JAX 生态里跑通了第一个神经网络。JAX 的门槛其实没那么高但它带来的性能收益特别是在大规模训练场景下绝对值得投入时间去学。作者Dr. Ashish Bamania