扩展# DTensor 概念 概述 安装 DTensor 的分布式张量模型 网格 布局 单客户端和多客户端应用 DTensor 作为分片张量 DTensor 剖析 布局和 fetch_layout 张量分量、pack 和 unpack 将 DTensor 分片到网格 二维网格上的完全分片 2 秩张量 二维网格上的完全复制 2 秩张量 二维网格上的混合 2 秩张量 Tensor.numpy() 和分片 DTensor DTensor 的相关 TensorFlow API DTensor 作为运算对象 完全复制输入和输出 沿收缩轴分片运算对象 额外分片 DTensor 作为输出 发出单个 TensorFlow 运算的 API 发出多个 TensorFlow 运算的 API 从 tf.Variable 到 dtensor.DVariable 后续步骤 高效的 TensorFlow 2 概述 设置 惯用 TensorFlow 2 的建议 将代码重构为更小的模块 调整某些 tf.keras.optimizer 的默认学习率 使用 tf.Module 和 Keras 层管理变量 结合 tf.data.Dataset 和 tf.function 使用 Keras 训练循环 自定义训练并编写自己的循环 通过 Python 控制流充分利用 tf.function 新型指标和损失 使用指标收集和显示数据 Keras 指标名称 调试 不要在您的对象中保留 tf.Tensors 资源和延伸阅读 扩展程序类型 安装 扩展程序类型 支持的 API 要求 字段类型 可变性 ExtensionType 添加的功能 构造函数 可打印表示 相等运算符 验证方法 强制不变性 嵌套 TypeSpec 自定义 ExtensionType 重写默认可打印表示 定义方法 定义类方法和静态方法 定义属性 重写默认构造函数 重写默认相等运算符 (__eq__) 使用前向引用 定义子类 定义私有字段 自定义 ExtensionType 的 TypeSpec 张量 API 调度 单个 API 的调度 所有一元逐元素 API 的调度 所有二进制逐元素 API 的调度 可批处理 ExtensionType BatchableExtensionType 示例:Network 支持 ExtensionType 的 TensorFlow API @tf.function 具体函数 控制流运算 Autograph 控制流 Keras Keras 示例:Network Keras 示例:MaskedTensor SavedModel 示例:保存 Keras 模型 示例:保存自定义模型 在 ExtensionType 不可用时加载 SavedModel ExtensionType 与 TensorFlow Serving 一起使用 数据集 使用扩展程序类型构建数据集 使用扩展程序类型批处理和取消批处理数据集 使用 JAX2TF 导入 JAX 模型 安装 下载并准备 MNIST 数据集 配置训练 使用 Flax 创建模型 编写训练步骤函数 编写训练循环 创建模型和优化器(使用 Optax) 训练模型 部分训练模型 保存推断所需的内容 保存一切 重新加载模型 在 TensorFlow 中继续训练转换后的 JAX 模型 后续步骤 Keras:TensorFlow 的高级 API 哪些人应当使用 Keras Keras API 组件 层 模型 其他 API 和工具 后续步骤 自动将代码升级到 TensorFlow 2 兼容性模块 推荐的升级流程 使用升级脚本 设置 读取帮助 TF1 代码示例 单个文件 目录树 详细报告 “安全”模式 注意事项