mirror of https://github.com/hpcaitech/ColossalAI
41 lines
1.2 KiB
Python
41 lines
1.2 KiB
Python
|
import os
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
BATCH_SIZE = 128
|
||
|
|
IMG_SIZE = 224
|
||
|
|
DIM = 768
|
||
|
|
NUM_CLASSES = 10
|
||
|
|
NUM_ATTN_HEADS = 12
|
||
|
|
|
||
|
|
# resnet 18
|
||
|
|
model = dict(type='VanillaResNet',
|
||
|
|
block_type='ResNetBasicBlock',
|
||
|
|
layers=[2, 2, 2, 2],
|
||
|
|
num_cls=10)
|
||
|
|
|
||
|
|
parallel = dict(
|
||
|
|
pipeline=dict(size=1),
|
||
|
|
tensor=dict(size=1, mode=None)
|
||
|
|
)
|
||
|
|
|
||
|
|
train_data = dict(dataset=dict(type='CIFAR10Dataset',
|
||
|
|
root=Path(os.environ['DATA']),
|
||
|
|
download=True,
|
||
|
|
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,
|
||
|
|
drop_last=True))
|
||
|
|
|
||
|
|
optimizer = dict(type='Adam', lr=0.001)
|
||
|
|
|
||
|
|
loss = dict(type='CrossEntropyLoss')
|
||
|
|
|