mirror of https://github.com/hpcaitech/ColossalAI
58 lines
1.1 KiB
ReStructuredText
58 lines
1.1 KiB
ReStructuredText
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:
|