colossalai.nn.layer.vanilla.layers

colossalai.nn.layer.vanilla.layers.drop_path(x, drop_prob=0.0, training=False)

Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). This is the same as the DropConnect impl I created for EfficientNet, etc networks, however, the original name is misleading as ‘Drop Connect’ is a different form of dropout in a separate paper… See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 … I’ve opted for changing the layer and argument names to ‘drop path’ rather than mix DropConnect as a layer name and use ‘survival rate’ as the argument.

class colossalai.nn.layer.vanilla.layers.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.layers.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.layers.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.

class colossalai.nn.layer.vanilla.layers.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.layers.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