mirror of https://github.com/hpcaitech/ColossalAI
92 lines
1.7 KiB
Python
92 lines
1.7 KiB
Python
#!/usr/bin/env python
|
|
# -*- encoding: utf-8 -*-
|
|
import os
|
|
from pathlib import Path
|
|
|
|
BATCH_SIZE = 128
|
|
IMG_SIZE = 224
|
|
NUM_CLS = 1000
|
|
|
|
# resnet 18
|
|
model = dict(
|
|
type='VanillaResNet',
|
|
block_type='ResNetBottleneck',
|
|
layers=[3, 4, 6, 3],
|
|
num_cls=NUM_CLS
|
|
)
|
|
|
|
train_data = dict(
|
|
dataset=dict(
|
|
type='CIFAR10Dataset',
|
|
root=Path(os.environ['DATA']),
|
|
transform_pipeline=[
|
|
dict(type='RandomResizedCrop', size=IMG_SIZE),
|
|
dict(type='RandomHorizontalFlip'),
|
|
dict(type='ToTensor'),
|
|
dict(type='Normalize', mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
|
|
]
|
|
),
|
|
dataloader=dict(
|
|
batch_size=64,
|
|
pin_memory=True,
|
|
num_workers=4,
|
|
sampler=dict(
|
|
type='DataParallelSampler',
|
|
shuffle=True,
|
|
)
|
|
)
|
|
)
|
|
|
|
test_data = dict(
|
|
dataset=dict(
|
|
type='CIFAR10Dataset',
|
|
root=Path(os.environ['DATA']),
|
|
train=False,
|
|
transform_pipeline=[
|
|
dict(type='Resize', size=(IMG_SIZE, IMG_SIZE)),
|
|
dict(type='ToTensor'),
|
|
dict(type='Normalize', mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
|
|
]
|
|
),
|
|
dataloader=dict(
|
|
batch_size=BATCH_SIZE,
|
|
pin_memory=True,
|
|
num_workers=4,
|
|
)
|
|
)
|
|
|
|
dist_initializer = [
|
|
dict(type='DataParallelInitializer'),
|
|
]
|
|
|
|
parallelization = dict(
|
|
pipeline=1,
|
|
tensor=1,
|
|
sequence=-1
|
|
)
|
|
|
|
optimizer = dict(
|
|
type='Adam',
|
|
lr=0.01
|
|
)
|
|
|
|
loss = dict(
|
|
type='CrossEntropyLoss'
|
|
)
|
|
|
|
trainer = dict(
|
|
max_epochs=5,
|
|
max_iters=1000
|
|
)
|
|
|
|
amp = dict(
|
|
fp16=None,
|
|
)
|
|
|
|
level = 2
|
|
|
|
parallel = dict(
|
|
pipeline=dict(size=1),
|
|
tensor=dict(size=1, mode=None)
|
|
)
|