google-deepmind/sonnet

TensorFlow-based neural network library

What it solves

Sonnet 是一個機器學習研究的函式庫，簡化了神經網路的建構。它提供一套可組合的抽象，讓研究人員能在不被限制於特定訓練框架或固定工作流程的情況下，構建複雜模型。

How it works

此函式庫建立於 TensorFlow 2 之上，核心概念為 snt.Module。Modules 是自包含、解耦的單元，可保存參數、其他 modules，以及處理輸入的方法。使用者可以使用預先定義好的 modules，如 snt.Linear 或 snt.nets.MLP，或透過繼承 snt.Module 自行建立。

Who it’s for

此函式庫設計給需要對模型與訓練流程擁有高度控制的機器學習研究者，特別是使用 TensorFlow 生態系的使用者。

Highlights

  • Unopinionated Design: 不包含訓練框架，訓練迴圈與 optimizer 的選擇交由使用者自行決定。
  • Composable Abstractions: 使用 snt.Module 系統，讓網路元件的嵌套變得簡單。
  • Flexible Serialization: 支援 Python 的 pickle、TensorFlow Checkpointing（保存訓練進度）以及 TensorFlow Saved Model（將模型導出為與 Python 原始碼解耦的格式）。
  • Distributed Training Support: 提供透過自訂 TensorFlow 分散策略進行分散式訓練的工具，同時保留使用者對梯度平均與 batch normalization 統計的控制。