Skip to main content
Version: v0.1.9

Initialize Features

Author: Shenggui Li, Siqi Mai

Prerequisite:

Introduction​

In this tutorial, we will cover the use of colossalai.initialize which injects features into your training components (e.g. model, optimizer, dataloader) seamlessly. Calling colossalai.initialize is the standard procedure before you run into your training loops.

In the section below, I will cover how colossalai.initialize works and what we should take note of.

Usage​

In a typical workflow, we will launch distributed environment at the beginning of our training script. Afterwards, we will instantiate our objects such as model, optimizer, loss function, dataloader etc. At this moment, colossalai.initialize can come in to inject features into these objects. A pseudo-code example is like below:

import colossalai
import torch
...


# launch distributed environment
colossalai.launch(config='./config.py', ...)

# create your objects
model = MyModel()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = torch.nn.CrossEntropyLoss()
train_dataloader = MyTrainDataloader()
test_dataloader = MyTrainDataloader()

# initialize features
engine, train_dataloader, test_dataloader, _ = colossalai.initialize(model,
optimizer,
criterion,
train_dataloader,
test_dataloader)

The colossalai.initialize function will return an Engine object. The engine object is a wrapper for model, optimizer and loss function. The engine object will run with features specified in the config file. More details about the engine can be found in the Use Engine and Trainer in Training.