Package: dks
Official Python package for Deep Kernel Shaping (DKS) and Tailored Activation Transformations (TAT)This Python package implements the activation function transformations, weight
initializations, and dataset preprocessing used in Deep Kernel Shaping (DKS) and
Tailored Activation Transformations (TAT). DKS and TAT, which were introduced in
theDKS paperandTAT paper, are methods for constructing/transforming
neural networks to make them much easier to train. For example, these methods
can be used in conjunction with K-FAC to train deep vanilla deep convnets
(without skip connections or normalization layers) as fast as standard ResNets
of the same depth.The package supports the JAX, PyTorch, and TensorFlow tensor programming
frameworks.Questions/comments about the code can be sent todks-dev@google.com.NOTE:we are not taking code contributions from Github at this time. All PRs
from Github will be rejected. Instead, please email us if you find a bug.UsageFor each of the supported tensor programming frameworks, there is a
corresponding subpackage which handles the activation function transformations,
weight initializations, and (optional) data preprocessing. (These aredks.jax,dks.pytorch, anddks.tensorflow.) It's up to the user to import these and
use them appropriately within their model code. Activation functions are
transformed by the functionget_transformed_activations()in the moduleactivation_transformof the appropriate subpackage. Sampling initial
parameters is done using functions inside of the moduleparameter_sampling_functionsof said subpackage. And data preprocessing is
done using the functionper_location_normalizationinside of the moduledata_preprocessingof said subpackage. Note that in order to avoid having to
import all of the tensor programming frameworks, the user is required to
individually import whatever framework subpackage they want. e.g.import dks.jax. Meanwhile,import dkswon't actually do anything.get_transformed_activations()requires the user to pass either the "maximal
slope function" for DKS, the "subnet maximizing function" for TAT with Leaky
ReLUs, or the "maximal curvature function" for TAT with smooth activation
functions. (The subnet maximizing function also handles DKS and TAT with smooth
activations.) These are special functions that encode information about the
particular model architecture. See the section titled "Summary of our method" of
theDKS paperfor a procedure to construct the maximal slope function for a
given model, or the appendix section titled "Additional details and pseudocode
for activation function transformations" of theTAT paperfor procedures to
construct the other two functions.In addition to these things, the user is responsible for ensuring that their
model meets the architectural requirements of DKS/TAT, and for converting any
weighted sums into "normalized sums" (which are weighted sums whose
non-trainable weights have a sum of squares equal to 1). See the section titled
"Summary of our method" of theDKS paperfor more details.Note that the data preprocessing method implemented, called Per-Location
Normalization (PLN), may not always be needed in practice, but we have observed
certain situations where not using can lead to problems. (For example, training
on datasets that contain all-zero pixels, such as CIFAR-10.) Also
note that ReLUs are only partially supported by DKS, and unsupported by TAT, and
so their use ishighlydiscouraged. Instead, one should use Leaky ReLUs, which
are fully supported by DKS, and work especially well with TAT.Exampledks.examples.haiku.modified_resnetis aHaikuResNet model which has been
modified as described in the DKS/TAT papers, and includes support for both DKS
and TAT. When constructed with its default arguments, it removes the
normalization layers and skip connections found in standard ResNets, making it a
"vanilla network". It can be used as an instructive example for how to build
DKS/TAT models using this package. See the section titled "Application to
various modified ResNets" from theDKS paperfor more details.InstallationThis package can be installed directly from GitHub usingpipwithpipinstallgit+https://github.com/deepmind/dks.gitorpipinstall-egit+https://github.com/deepmind/dks.git#egg=dks[<extras>]Or from PyPI withpipinstalldksorpipinstalldks[<extras>]Here<extras>is a common-separated list of strings (with no spaces) that can
be passed to install extra dependencies for different tensor programming
frameworks. Valid strings arejax,tf, andpytorch. So for example, to
installdkswith the extra requirements for JAX and PyTorch, one doespipinstalldks[jax,pytorch]TestingTo run tests in a Python virtual environment with specific pinned versions of
all the dependencies one can do:gitclonehttps://github.com/deepmind/dks.gitcddks
./test.shHowever, it is strongly recommended that you run the tests in the same Python
environment (with the same package versions) as you plan to actually usedks.
This can be accomplished by installingdksfor all three tensors programming
frameworks (e.g. withpip install dks[jax,pytorch,tf]or some other
installation method), and then doingpipinstallpytest-xdist
gitclonehttps://github.com/deepmind/dks.gitcddks
python-mpytest-n16testsDisclaimerThis is not an official Google product.