mirror of https://github.com/hpcaitech/ColossalAI
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
21 lines
361 B
21 lines
361 B
import os
|
|
import model
|
|
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=4),
|
|
tensor=dict(size=1, mode=None)
|
|
)
|