kd.nn

[[Source]]

Collection of nn.Modules to build neural networks.

Symbols

Class

kd.nn.AddEmbedding	Helper Module for adding a PositionEmbedding e.g. in a knn.Sequential.
kd.nn.AddLearnedEmbedding	Adds learned positional embeddings to the inputs.
kd.nn.AttentionModule	Interface specification for Attention modules.
kd.nn.Dropout	Wrapper around nn.Dropout but using kd.nn.train_property.
kd.nn.DummyModel	Empty model that ignores inputs and always produces a single logit of 42.
kd.nn.ExternalModule	Module that is defined outside Kauldron.
kd.nn.FlatAutoencoder	Very simple auto-encoder class to showcase using keys and submodules.
kd.nn.FourierEmbedding	Apply Fourier position embedding to a grid of coordinates.
kd.nn.Identity	Module that applies the identity function to a single tensor.
kd.nn.ImageTokenizer	Interface for modules that convert images into tokens.
kd.nn.ImprovedMultiHeadDotProductAttention	Multi-head dot-product attention.
kd.nn.LearnedEmbedding	Learned positional embeddings.
kd.nn.MultiHeadDotProductAttention	Wrapper around nn.MultiHeadDotProductAttention using knn.train_property.
kd.nn.NormModule	Interface specification for norm modules (to be used as type annotation).
kd.nn.ParallelAttentionBlock	Parallel self attention (see Vit22B paper: arxiv.org/abs/2302.05442).
kd.nn.Patchify	Patchify an image, as in ViT (without linear embedding).
kd.nn.PatchifyEmbed	Patchify and linearly embed and image, as in ViT.
kd.nn.PostNormBlock	Post-LN Transformer layer (not recommended).
kd.nn.PreNormBlock	Pre-LN Transformer layer (default transformer layer).
kd.nn.Rearrange	Wrapper around einops.rearrange for usage e.g. in nn.Sequential.
kd.nn.Reduce	Wrapper around einops.reduce for usage e.g. in nn.Sequential.
kd.nn.Sequential	Like nn.Sequential but allows configuring input and output keys.
kd.nn.TransformerBlock	Interface definition for transformer blocks (for use in type annotations).
kd.nn.TransformerMLP	Simple MLP with a single hidden layer for use in Transformer blocks.
kd.nn.Vit	Basic Vision Transformer classifer with GAP.
kd.nn.VitEncoder	Basic Vit Encoder.
kd.nn.WrapperModule	Base class to wrapper a module.
kd.nn.ZeroEmbedding	Embedding that returns zero (for deactivating position embeddings).

Function

kd.nn.convert_to_fourier_features	Convert inputs to Fourier features, e.g. for positional encoding.
kd.nn.interms_property	interms property that makes storing intermediates more convenient.
kd.nn.set_train_property	Set the self.is_training state to the given value.
kd.nn.train_property	is_training property.