ColossalAI/docs/colossalai/colossalai.nn.rst

colossalai.nn
=============

*This part contains different colossalai layers for constructing your model.
You can easily use them as the way of using layers in torch.nn.*

*Now colossalai support layer types below:* ``Linear``, ``Classifier``, ``Embedding``,
``PatchEmbedding``, ``LayerNorm``, ``Dropout`` *for different parallelisms.*

.. toctree::
   :maxdepth: 2

   colossalai.nn.layer

*This part contains different loss functions for different parallelisms.*

.. toctree::
   :maxdepth: 2

   colossalai.nn.loss

*This part contains different learning rate schedules to control your learning rate
in training process*

.. toctree::
   :maxdepth: 2

   colossalai.nn.lr_scheduler

*This part contains different metric to measure performance of your model.*

.. toctree::
   :maxdepth: 2

   colossalai.nn.metric

*This part contains some colossalai optimizers*

.. toctree::
   :maxdepth: 2

   colossalai.nn.model

.. toctree::
   :maxdepth: 2

   colossalai.nn.optimizer

*This part contains different methods to initialize weights.*

.. toctree::
   :maxdepth: 2

   colossalai.nn.init

.. automodule:: colossalai.nn
   :members: