首页
学习
活动
专区
圈层
工具
发布
    • 综合排序
    • 最热优先
    • 最新优先
    时间不限
  • 来自专栏DeepHub IMBA

    JAX 核心特性详解:纯函数、JIT 编译、自动微分等十大必知概念

    Flax 代码风格接近 PyTorch,如果你用过 PyTorch 会很快上手。 Flax 提供了丰富的层和操作。 =256) flax.linen.MultiHeadDotProductAttention(num_heads=8) flax.linen.SelfAttention(num_heads=8 ) 归一化 支持多种方式: flax.linen.BatchNorm() flax.linen.LayerNorm() flax.linen.GroupNorm(num_groups=32 ) flax.linen.RMSNorm() 激活和 Dropout: flax.linen.relu(x) flax.linen.gelu(x) flax.linen.sigmoid (x) flax.linen.tanh(x) flax.linen.Dropout(rate=0.1) 池化: flax.linen.avg_pool(x, window_shape=(2,2

    33710编辑于 2025-11-26
  • 来自专栏DeepHub IMBA

    Google开源Tunix:JAX生态的LLM微调方案来了

    Tunix是什么 这是个构建在JAX之上的后训练库,和Flax NNX集成得比较紧密。 q git+https://github.com/google/tunix pip install -q git+https://github.com/google/qwix # Flax 需要升级到最新版 pip uninstall -q -y flax pip install -q git+https://github.com/google/flax.git 完整流程 第一步, 从Kaggle拉预训练checkpoint: import kagglehub model_path = "google/gemma/flax/2b" kaggle_ckpt_path = kagglehub.model_download(model_path) 初始化模型和tokenizer: from flax import nnx from tunix.models.gemma

    23210编辑于 2025-11-15
  • 来自专栏DeepHub IMBA

    深度解析 Google JAX 全栈:带你上手开发,从零构建神经网络

    2、Flax 基于 JAX 的神经网络训练库。Flax 的核心现在是 NNX (Neural Networks for JAX)。 之前有个 Flax Linen,是那种无状态、函数式风格的 API。 构建网络 用 Flax NNX 搭建一个带 SELU 激活函数的简单前馈网络。 from flax import nnx class DigitClassifier(nnx.Module): def __init__(self, n_features, n_hidden Flax NNX 把它俩封装成了装饰器 nnx.jit 和 nnx.grad,用起来更方便。 9、训练循环 跑 500 epoch,每 100 轮显示 Loss。

    26810编辑于 2025-12-19
  • 来自专栏数据派THU

    使用JAX实现完整的Vision Transformer

    来源:DeepHub IMBA 本文约3200字,建议阅读10+分钟 本文将展示如何使用JAX/Flax实现Vision Transformer (ViT),以及如何使用JAX/Flax训练ViT。 细节实现 下面,我们将使用JAX/Flax创建每个模块。 1. 图像到展平的图像补丁 下面的代码从输入图像中提取图像补丁。 (x, train) # MLP head x = x[:, 0] # [CLS] token x = self.cls_head(x) return x 使用JAX/Flax 训练 现在已经创建了模型,下面就是使用JAX/Flax来训练。 return model, params, main_rng vit_model, vit_params, vit_rng = initialize_model() 创建TrainState 在Flax

    72221编辑于 2023-03-29
  • 来自专栏数据科学(冷冻工厂)

    深度学习长文|使用 JAX 进行 AI 模型训练

    在本文中,我们将使用 Flax,据本文撰写时的观察,它似乎是最受欢迎的选择。 JAX 幕后花絮 — XLA 编译 JAX 的强大之处在于它利用了 XLA 编译技术。 import torch import jax, flax, optax import jax.numpy as jnp def get_model(use_jax=False): from JAX 训练循环依赖于 Flax TrainState 对象,其定义遵循在 Flax 中训练 ML 模型的基本教程: @jax.jit def train_step_jax(train_state, batch common-cu121-v20240514-ubuntu-2204-py310),并预装了 PyTorch(版本 2.3.0)、PyTorch/XLA(版本 2.3.0)、JAX(版本 0.4.28)、Flax 实际上,考虑到 PyTorch 背后庞大的开发者社区,以及 PyTorch 所原生支持而 JAX/FLAX 尚未涵盖的众多特性(例如自动混合精度、先进的注意力机制层,至少在本文撰写时),有人可能会强烈主张没有必要投入时间去掌握

    99211编辑于 2024-06-18
  • 来自专栏Elastic Stack专栏

    如何将本地transformer模型部署到Elasticsearch

    . ├── README.md ├── added_tokens.json ├── config.json ├── flax_model.msgpack ├── model.safetensors ├─ models--bert-base-uncased/snapshots/0a6aa9128b6194f4f3c4db429b6cb4891cdb421b but there is a file for Flax Use `from_flax=True` to load this model from those weights.

    4.9K31编辑于 2023-06-15
  • 来自专栏信数据得永生

    Transformers 4.37 中文文档(二十一)

    此模型还是一个 Flax Linen flax.nn.Module子类。将其用作常规 Flax 模块,并参考 Flax 文档以获取有关一般用法和行为的所有相关信息。 此模型还是 Flax Linen flax.nn.Module子类。将其用作常规 Flax 模块,并参考 Flax 文档以获取有关一般用法和行为的所有相关信息。 此模型还是一个 Flax 亚麻flax.nn.Module子类。将其用作常规 Flax 模块,并参考 Flax 文档以获取有关一般用法和行为的所有相关信息。 此模型还是 Flax 亚麻Linen子类。将其用作常规 Flax 模块,并参考 Flax 文档以获取有关一般用法和行为的所有相关信息。 这个模型也是一个 Flax Linen flax.nn.Module子类。将其用作常规 Flax 模块,并参考 Flax 文档以了解所有与一般用法和行为相关的事项。

    62210编辑于 2024-06-26
  • 来自专栏信数据得永生

    Transformers 4.37 中文文档(三十)

    这个模型也是一个 Flax Linen flax.nn.Module子类。将其用作常规的 Flax 模块,并参考 Flax 文档以了解与一般用法和行为相关的所有内容。 这个模型也是一个 Flax 亚麻flax.nn.Module子类。将其用作常规的 Flax 模块,并参考 Flax 文档以获取与一般用法和行为相关的所有内容。 这个模型也是 Flax Linen flax.nn.Module子类。将其用作常规 Flax 模块,并参考 Flax 文档以获取与一般用法和行为相关的所有内容。 此模型也是 Flax 亚麻flax.nn.Module子类。 将其用作常规 Flax 模块,并参考 Flax 文档以获取与一般用法和行为相关的所有内容。 此模型也是 Flax 亚麻flax.nn.Module子类。 将其用作常规 Flax 模块,并参考 Flax 文档以获取与一般用法和行为相关的所有内容。

    1.3K10编辑于 2024-08-26
  • 来自专栏信数据得永生

    Transformers 4.37 中文文档(二十三)

    这个模型也是一个flax.linen.Module子类。将其用作常规的 Flax 亚麻模块,并参考 Flax 文档以获取与一般用法和行为相关的所有内容。 此模型也是flax.linen.Module的子类。将其用作常规的 Flax linen 模块,并参考 Flax 文档以获取有关一般用法和行为的所有信息。 此模型还是flax.linen.Module子类。将其用作常规 Flax 亚麻模块,并参考 Flax 文档以获取有关一般用法和行为的所有相关信息。 此模型还是flax.linen.Module子类。将其用作常规的 Flax 亚麻模块,并参考 Flax 文档以获取有关一般用法和行为的所有相关信息。 此模型还是一个flax.linen.Module子类。将其用作常规的 Flax 亚麻模块,并参考 Flax 文档以获取有关一般用法和行为的所有相关信息。

    48910编辑于 2024-06-26
  • 来自专栏数据猿

    日本机器自动化初创企业Cinnamon获天使轮融资

    公司的发展重点是RPA(机器人过程自动化)第2类解决方案Flax扫描仪,这一解决方案使用公司自己研发的文档阅读引擎“Cinnamon AI”。 除了Flax扫描仪,Cinnamon还提供其他人工智能解决方案,如匹配用户和产品的Lapis引擎以及聊天机器人Scuro Bot。

    91590发布于 2018-04-25
  • 来自专栏皮皮鲁的AI星球

    TensorFlow被废了,谷歌家的新王储JAX到底是啥?

    于是有了DeepMind的 haiku ,Google的 flax,和其他各种各样的库。 JAX是纯函数的,代码写起来和tf、torch也不太一样。 不过,flax 和 haiku 也有不少市场了。大名鼎鼎的AlphaFold就是用 haiku 写的。 但大家都在学JAX JAX到底好不好我不敢说。但是大家都在学它。

    1.1K10编辑于 2022-09-20
  • 来自专栏信数据得永生

    Transformers 4.37 中文文档(四十三)

    该模型还是 Flax Linen 的flax.nn.Module子类。将其用作常规 Flax 模块,并参考 Flax 文档以获取有关一般用法和行为的所有相关信息。 此模型还是 Flax Linen flax.nn.Module子类。将其用作常规 Flax 模块,并参考 Flax 文档以获取有关一般用法和行为的所有相关信息。 返回 transformers.modeling_flax_outputs.FlaxSeq2SeqLMOutput 或tuple(torch.FloatTensor) transformers.modeling_flax_outputs.FlaxSeq2SeqLMOutput 此模型还是 Flax 亚麻flax.nn.Module子类。将其用作常规 Flax 模块,并参考 Flax 文档以获取有关一般用法和行为的所有相关信息。 此模型还是 Flax 亚麻flax.nn.Module子类。将其用作常规 Flax 模块,并参考 Flax 文档以获取有关一般用法和行为的所有相关信息。

    65410编辑于 2024-06-26
  • 来自专栏信数据得永生

    Transformers 4.37 中文文档(二十)

    此模型还是一个flax.linen.Module子类。将其用作常规的 Flax linen 模块,并参考 Flax 文档以获取与一般用法和行为相关的所有内容。 此模型还是一个flax.linen.Module子类。将其用作常规的 Flax 亚麻模块,并参考 Flax 文档以获取与一般用法和行为相关的所有事项。 此模型还是一个flax.linen.Module子类。将其用作常规的 Flax linen 模块,并参考 Flax 文档以获取与一般用法和行为相关的所有内容。 此模型还是一个flax.linen.Module子类。将其用作常规的 Flax linen 模块,并参考 Flax 文档以了解与一般用法和行为相关的所有事项。 此模型还是一个flax.linen.Module子类。将其用作常规的 Flax 亚麻模块,并参考 Flax 文档以了解所有与一般用法和行为相关的事项。

    97310编辑于 2024-06-26
  • 来自专栏信数据得永生

    Transformers 4.37 中文文档(二十五)

    此模型还是 Flax Linen flax.nn.Module子类。将其用作常规 Flax 模块,并参考 Flax 文档以获取有关一般用法和行为的所有相关信息。 返回 transformers.modeling_flax_outputs.FlaxBaseModelOutput 或tuple(torch.FloatTensor) 一个 transformers.modeling_flax_outputs.FlaxBaseModelOutput 此模型也是 Flax 亚麻flax.nn.Module子类。将其用作常规 Flax 模块,并参考 Flax 文档以获取与一般用法和行为相关的所有内容。 此模型还是 Flax 亚麻flax.nn.Module子类。将其用作常规 Flax 模块,并参考 Flax 文档以获取有关一般用法和行为的所有相关信息。 这个模型也是一个 Flax 亚麻flax.nn.Module子类。将其用作常规的 Flax 模块,并参考 Flax 文档以获取有关一般用法和行为的所有信息。

    46110编辑于 2024-06-26
  • 农作物害虫检测数据集VOC+YOLO格式18975张97类别

    brown plant hopper","cabbage army worm","cerodonta denticornis","corn borer","english grain aphid","flax army worm 框数 = 266 cerodonta denticornis 框数 = 27 corn borer 框数 = 440 english grain aphid 框数 = 310 flax

    23300编辑于 2025-07-20
  • 来自专栏信数据得永生

    Transformers 4.37 中文文档(五十六)

    此模型还是flax.linen.Module子类。将其用作常规 Flax 亚麻模块,并参考 Flax 文档以了解与一般用法和行为相关的所有事项。 此模型还是一个flax.linen.Module子类。将其用作常规 Flax linen 模块,并参考 Flax 文档以了解与一般使用和行为相关的所有事项。 这个模型也是一个flax.linen.Module子类。将其用作常规的 Flax 亚麻模块,并参考 Flax 文档以获取与一般用法和行为相关的所有内容。 此模型也是flax.linen.Module的子类。将其用作常规的 Flax 亚麻模块,并参考 Flax 文档以获取与一般用法和行为相关的所有内容。 此模型也是一个 flax.linen.Module 子类。将其用作常规的 Flax linen 模块,并参考 Flax 文档以了解与一般使用和行为相关的所有事项。

    49410编辑于 2024-06-26
  • 来自专栏HelloGitHub

    国产的开源数据库——GitHub 热点速览 Vol.52

    2.5 3D 游戏引擎:FlaxEngine 本周 star 增长数:750+ Flax Engine 是一个用 C++ 和 C# 编写的高质量的现代 3D 游戏引擎。 从惊艳的图形到强大的脚本,Flax 可以给你提供关于游戏的一切。 GitHub 地址→https://github.com/FlaxEngine/FlaxEngine ?

    92620发布于 2021-05-14
  • 来自专栏信数据得永生

    Transformers 4.37 中文文档(五十五)

    此模型还是flax.linen.Module的子类。将其用作常规的 Flax 亚麻模块,并参考 Flax 文档以获取与一般用法和行为相关的所有内容。 这个模型也是一个flax.linen.Module子类。将其用作常规的 Flax 亚麻模块,并参考 Flax 文档以获取与一般用法和行为相关的所有内容。 此模型也是一个flax.linen.Module子类。将其用作常规的 Flax linen 模块,并参考 Flax 文档以获取有关一般用法和行为的所有相关信息。 此模型还是 flax.linen.Module 的子类。将其用作常规的 Flax linen 模块,并参考 Flax 文档以了解与一般用法和行为相关的所有事项。 这个模型也是一个flax.linen.Module子类。将其用作常规的 Flax 亚麻模块,并参考 Flax 文档以获取与一般用法和行为相关的所有内容。

    64410编辑于 2024-06-26
  • 来自专栏HelloGitHub

    什么是 x10 开发工具?「GitHub 热点速览」

    GitHub Trending 周榜 2.1 3D 游戏引擎:FlaxEngine 本周 star 增长数 1,050+,主语言:C++、C# Flax Engine 是一个用 C++ 和 C# 编写的高质量 从令人惊叹的图形到功能强大的脚本,Flax 可为你的游戏提供一切。该引擎专为快速工作流程而设计,有非常多即开即用的功能。

    52940编辑于 2023-09-24
  • 来自专栏DeepHub IMBA

    JAX 训练加速指南:8 个让 TPU 满跑的工程实战习惯

    .prefetch(tf.data.AUTOTUNE)) # Convert to NumPy and prefetch onto devices from flax.jax_utils from flax import linen as nn class DeepBlock(nn.Module): @nn.compact 极简 TPU 训练示例 这基本包含了上面所有的内容 # Pseudo-skeleton (Flax + JAX + TPU) mesh = Mesh(np.array(jax.devices

    18210编辑于 2025-12-24
领券