mirror of https://github.com/hpcaitech/ColossalAI
22 lines
383 B
Python
22 lines
383 B
Python
import os
|
|
import model
|
|
from pathlib import Path
|
|
|
|
BATCH_SIZE = 128
|
|
IMG_SIZE = 224
|
|
DIM = 768
|
|
NUM_CLASSES = 10
|
|
NUM_ATTN_HEADS = 12
|
|
NUM_MICRO_BATCHES = 2
|
|
|
|
# resnet 18
|
|
model = dict(type='VanillaResNet',
|
|
block_type='ResNetBasicBlock',
|
|
layers=[2, 2, 2, 2],
|
|
num_cls=10)
|
|
|
|
parallel = dict(
|
|
pipeline=dict(size=4),
|
|
tensor=dict(size=1, mode=None)
|
|
)
|