/
githubmirror
/
ColossalAI
Обзор
Документация
Войти
/
githubmirror
/
ColossalAI
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
main
tests/kit/model_zoo/__init__.py
33 строки
1 KB
Edenzzzz
[Feature] Zigzag Ring attention (#5905)
16 авг 2024, 08:56
Не верифицирован
16 авг 2024, 08:56
f5c84af
Код
Авторство
О чём код?
import os from . import custom, diffusers, timm, torchaudio, torchvision, transformers from .executor import run_fwd, run_fwd_bwd from .registry import model_zoo # We pick a subset of models for fast testing in order to reduce the total testing time COMMON_MODELS = [ "custom_hanging_param_model", "custom_nested_model", "custom_repeated_computed_layers", "custom_simple_net", "diffusers_clip_text_model", "diffusers_auto_encoder_kl", "diffusers_unet2d_model", "timm_densenet", "timm_resnet", "timm_swin_transformer", "torchaudio_wav2vec2_base", "torchaudio_conformer", "transformers_bert_for_masked_lm", "transformers_bloom_for_causal_lm", "transformers_falcon_for_causal_lm", "transformers_chatglm_for_conditional_generation", "transformers_llama_for_causal_lm", "transformers_vit_for_masked_image_modeling", "transformers_mistral_for_causal_lm", ] IS_FAST_TEST = os.environ.get("FAST_TEST", "0") == "1" __all__ = ["model_zoo", "run_fwd", "run_fwd_bwd", "COMMON_MODELS", "IS_FAST_TEST"]