sonnet

Sonnet is an open-source neural network library developed by researchers at DeepMind, designed to sit on top of TensorFlow 2. Its primary goal is to provide simple, composable abstractions that facilitate machine learning research across various domains, including supervised/unsupervised learning and reinforcement learning. The core programming model in Sonnet revolves around the `snt.Module` class. Modules are self-contained units that can hold references to parameters, other modules, and methods. This design ensures that modules are decoupled from one another, allowing for high flexibility in how they are used. Unlike some frameworks, Sonnet does not enforce a specific training loop or opinionated structure; instead, it encourages users to build their own training logic or adopt existing solutions. Key features include: - **Predefined Modules**: Sonnet ships with common building blocks such as `snt.Linear`, `snt.Conv2D`, and `snt.BatchNorm`, as well as pre-built networks like MLPs. - **Custom Modules**: Users are encouraged to subclass `snt.Module` to create custom layers or networks. Subclassing provides automatic benefits like parameter management (`variables` and `trainable_variables` properties), name scoping for TensorBoard visualization, and support for TensorFlow checkpointing. - **Serialization**: The library supports saving and loading models via Python's `pickle` (though discouraged due to brittleness) and TensorFlow's native checkpointing and Saved Model formats. This allows for robust model persistence and deployment. - **Distributed Training**: Sonnet offers support for distributed training using custom TensorFlow distribution strategies. It gives users full control over gradient averaging and batch normalization statistics synchronization, rather than baking these behaviors into the library automatically. Installation requires TensorFlow 2 and TensorFlow Probability, followed by installing `dm-sonnet`. Examples are available via Google Colab notebooks for tasks like MNIST classification, GAN training, and distributed CIFAR-10 training.