mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-05-11 09:57:26 +00:00
0fedef4f3c
* integrated parallel layers for ease of building models * integrated 2.5d layers * cleaned codes and unit tests * added log metric by step hook; updated imagenet benchmark; fixed some bugs * reworked initialization; cleaned codes Co-authored-by: BoxiangW <45734921+BoxiangW@users.noreply.github.com>
20 lines
392 B
Python
20 lines
392 B
Python
|
|
import sys
|
|
from pathlib import Path
|
|
repo_path = Path(__file__).absolute().parents[2]
|
|
sys.path.append(str(repo_path))
|
|
|
|
try:
|
|
import model_zoo.vit.vision_transformer_from_config
|
|
except ImportError:
|
|
raise ImportError("model_zoo is not found, please check your path")
|
|
|
|
BATCH_SIZE = 8
|
|
IMG_SIZE = 32
|
|
PATCH_SIZE = 4
|
|
DIM = 512
|
|
NUM_ATTENTION_HEADS = 8
|
|
SUMMA_DIM = 2
|
|
NUM_CLASSES = 10
|
|
DEPTH = 6
|