colossalai.nn.layer.vanilla

class colossalai.nn.layer.vanilla.VanillaPatchEmbedding(img_size, patch_size, in_chans, embed_size, flatten=True, dtype=None, weight_initializer=<function kaiming_uniform_.<locals>.initializer>, bias_initializer=<function xavier_uniform_.<locals>.initializer>, position_embed_initializer=<function zeros_.<locals>.initializer>)

2D Image to Patch Embedding

Parameters

  • img_size (int) – image size

  • patch_size (int) – patch size

  • in_chans (int) – number of channels of input image

  • embed_size (int) – size of embedding

  • dtype (torch.dtype, optional) – The dtype of parameters, defaults to None

  • flatten (bool, optional) – whether to flatten output tensor, defaults to True

  • weight_initializer (Callable, optional) – The intializer of weight, defaults to kaiming uniform initializer

  • bias_initializer (Callable, optional) – The intializer of bias, defaults to xavier uniform initializer

  • position_embed_initializer (Callable, optional) – The intializer of position embedding, defaults to zero

class colossalai.nn.layer.vanilla.VanillaClassifier(in_features, num_classes, weight=None, bias=True, dtype=None, weight_initializer=<function kaiming_uniform_.<locals>.initializer>, bias_initializer=<function xavier_uniform_.<locals>.initializer>)

Dense linear classifier

Parameters

  • in_features (int) – size of each input sample

  • num_classes (int) – number of classes

  • weight (torch.nn.Parameter, optional) – weight of the classifier, defaults to True

  • bias (bool, optional) – If set to False, the layer will not learn an additive bias, defaults to True

  • dtype (torch.dtype, optional) – The dtype of parameters, defaults to None

  • weight_initializer (Callable, optional) – The intializer of weight, defaults to kaiming uniform initializer

  • bias_initializer (Callable, optional) – The intializer of bias, defaults to xavier uniform initializer

class colossalai.nn.layer.vanilla.DropPath(drop_prob=None)

Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). Adapted from https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/layers/drop.py

class colossalai.nn.layer.vanilla.WrappedDropout(p=0.5, inplace=False, mode=None)

Same as torch.nn.Dropout. But it is wrapped with the context of seed manager.

class colossalai.nn.layer.vanilla.WrappedDropPath(p=0.0, mode=None)

Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). Here, it is wrapped with the context of seed manager.