colossalai.nn.layer.colossalai_layer.embedding

class colossalai.nn.layer.colossalai_layer.embedding.Embedding(num_embeddings, embedding_dim, padding_idx=None, dtype=None, weight_initializer=<function normal_.<locals>.initializer>, vocab_parallel_limit=2048, *args, **kwargs)

Embedding for colossalai

Parameters

• num_embeddings (int) – number of embeddings

• embedding_dim (int) – dimension of embedding

• padding_idx (int, optional) – index of padding, defaults to None

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

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

• args – Args used in F.embedding

• kwargs – Kwargs used in F.embedding

class colossalai.nn.layer.colossalai_layer.embedding.PatchEmbedding(img_size, patch_size, in_chans, embed_size, dtype=None, flatten=True, 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