<jupyter_start><jupyter_text>Train an image classifier<jupyter_code>!pip install ColossalAI deepspeed
import colossalai
from colossalai.engine import Engine, NoPipelineSchedule
from colossalai.trainer import Trainer
from colossalai.context import Config
import torch<jupyter_output>Please install apex to use FP16 Optimizer
Apex should be installed to use the FP16 optimizer
apex is required for mixed precision training<jupyter_text>First, we should initialize distributed environment. Though we just use single GPU in this example, we still need initialize distributed environment for compatibility. We just consider the simplest case here, so we just set the number of parallel processes to 1.<jupyter_code>parallel_cfg = Config(dict(parallel=dict(
    data=dict(size=1),
    pipeline=dict(size=1),
    tensor=dict(size=1, mode=None),
)))
colossalai.init_dist(config=parallel_cfg,
          local_rank=0,
          world_size=1,
          host='127.0.0.1',
          port=8888,
          backend='nccl')<jupyter_output>colossalai - torch.distributed.distributed_c10d - 2021-10-15 03:27:51,596 INFO: Added key: store_based_barrier_key:1 to store for rank: 0
colossalai - torch.distributed.distributed_c10d - 2021-10-15 03:27:51,598 INFO: Rank 0: Completed store-based barrier for 1 nodes.
colossalai - torch.distributed.distributed_c10d - 2021-10-15 03:27:51,602 INFO: Added key: store_based_barrier_key:2 to store for rank: 0
colossalai - torch.distributed.distributed_c10d - 2021-10-15 03:27:51,605 INFO: Rank 0: Completed store-based barrier for 1 nodes.
colossalai - torch.distributed.distributed_c10d - 2021-10-15 03:27:51,608 INFO: Added key: store_based_barrier_key:3 to store for rank: 0
colossalai - torch.distributed.distributed_c10d - 2021-10-15 03:27:51,610 INFO: Rank 0: Completed store-based barrier for 1 nodes.<jupyter_text>Load and normalize the CIFAR10 training and test datasets using `colossalai.nn.data`. Note that we have wrapped `torchvision.transforms`, so that we can simply use the config dict to use them.<jupyter_code>transform_cfg = [
    dict(type='ToTensor'),
    dict(type='Normalize',
        mean=[0.4914, 0.4822, 0.4465],
        std=[0.2023, 0.1994, 0.2010]),
]

batch_size = 128

trainset = colossalai.nn.data.CIFAR10Dataset(transform_cfg, root='./data', train=True)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=2)

testset = colossalai.nn.data.CIFAR10Dataset(transform_cfg, root='./data', train=False)
testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size, shuffle=False, num_workers=2)<jupyter_output>Files already downloaded and verified
Files already downloaded and verified<jupyter_text>We just define a simple Convolutional Neural Network here.<jupyter_code>import torch.nn as nn
import torch.nn.functional as F


class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = torch.flatten(x, 1) # flatten all dimensions except batch
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


model = Net().cuda()<jupyter_output><empty_output><jupyter_text>Define a Loss function and optimizer. And then we use them to initialize `Engine` and `Trainer`. We provide various training / evaluating hooks. In this case, we just use the simplest hooks which can compute and print loss and accuracy.<jupyter_code>import torch.optim as optim

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
schedule = NoPipelineSchedule()
engine = Engine(
        model=model,
        criterion=criterion,
        optimizer=optimizer,
        lr_scheduler=None,
        schedule=schedule
    )
trainer = Trainer(engine=engine,
          hooks_cfg=[dict(type='LossHook'), dict(type='LogMetricByEpochHook'), dict(type='AccuracyHook')],
          verbose=True)<jupyter_output>colossalai - rank_0 - 2021-10-15 03:27:56,018 WARNING: No gradient handler is set up, please make sure you do not need to all-reduce the gradients after a training step.
colossalai - rank_0 - 2021-10-15 03:27:56,024 INFO: build LogMetricByEpochHook for train, priority = 1
colossalai - rank_0 - 2021-10-15 03:27:56,026 INFO: build LossHook for train, priority = 10
colossalai - rank_0 - 2021-10-15 03:27:56,029 INFO: build AccuracyHook for train, priority = 10<jupyter_text>Then we set training configs. We train our model for 10 epochs and it will be evaluated every 1 epoch. Set `display_progress` to `True` to display the training / evaluating progress bar.<jupyter_code>num_epochs = 10
test_interval = 1
trainer.fit(
        train_dataloader=trainloader,
        test_dataloader=testloader,
        max_epochs=num_epochs,
        display_progress=True,
        test_interval=test_interval
    )<jupyter_output>[Epoch 0 train]:   0%|          | 0/391 [00:00<?, ?it/s]/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py:718: UserWarning: Named tensors and all their associated APIs are an experimental feature and subject to change. Please do not use them for anything important until they are released as stable. (Triggered internally at  /pytorch/c10/core/TensorImpl.h:1156.)
  return torch.max_pool2d(input, kernel_size, stride, padding, dilation, ceil_mode)
[Epoch 0 train]: 100%|██████████| 391/391 [00:14<00:00, 26.82it/s]
colossalai - rank_0 - 2021-10-15 03:28:11,088 INFO: Training - Epoch 1 - LogMetricByEpochHook: Loss = 2.29158
[Epoch 0 val]: 100%|██████████| 79/79 [00:02<00:00, 28.66it/s]
colossalai - rank_0 - 2021-10-15 03:28:14,040 INFO: Testing - Epoch 1 - LogMetricByEpochHook: Loss = 2.26517, Accuracy = 0.14820
[Epoch 1 train]: 100%|██████████| 391/391 [00:14<00:00, 26.31it/s]
colossalai - rank_0 - 2021-10-15 03:28:29,059 INFO: Training - Epoch 2 - LogMetricByEpochHook: Loss = 2.[...]
{'fraction_non_alphanumeric': 0.11087696547252392, 'fraction_numerical': 0.07732209434268114, 'mean_word_length': 4.219966159052453, 'pattern_counts': {'":': 0, '<': 24, '<?xml version=': 0, '>': 20, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 1, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '99509', 'license': 'Apache-2.0', 'n_tokens_mistral': 2430, 'n_tokens_neox': 2058, 'n_words': 569, 'path': 'examples/colossal_cifar_demo.ipynb', 'repo_name': 'rahulgupta9202/ColossalAI', 'size': '6169'}