github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-06 03:20:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

[test] added transformers models to test model zoo (#3135)

Browse Source
This commit is contained in:
Frank Lee
2023-03-15 11:26:10 +08:00
committed by GitHub
parent a674c63348
commit 6d48eb0560
12 changed files with 339 additions and 193 deletions
Download Patch File Download Diff File Expand all files Collapse all files

49
tests/kit/model_zoo/transformers/gpt.py Normal file
View File

@@ -0,0 +1,49 @@
import torch
import transformers
from ..registry import ModelAttribute, model_zoo
# ===============================
# Register single-sentence GPT
# ===============================
BATCH_SIZE = 2
SEQ_LENGTH = 16
def data_gen():
input_ids = torch.zeros((BATCH_SIZE, SEQ_LENGTH), dtype=torch.int64)
token_type_ids = torch.zeros((BATCH_SIZE, SEQ_LENGTH), dtype=torch.int64)
attention_mask = torch.zeros((BATCH_SIZE, SEQ_LENGTH), dtype=torch.int64)
return dict(input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask)
output_transform_fn = lambda x: x
config = transformers.GPT2Config(n_position=64, n_layer=2, n_head=4)
# register the following models
model_zoo.register(name='transformers_gpt',
model_fn=lambda: transformers.GPT2Model(config),
data_gen_fn=data_gen,
output_transform_fn=output_transform_fn,
model_attribute=ModelAttribute(has_control_flow=True))
model_zoo.register(name='transformers_gpt_lm',
model_fn=lambda: transformers.GPT2LMHeadModel(config),
data_gen_fn=data_gen,
output_transform_fn=output_transform_fn,
model_attribute=ModelAttribute(has_control_flow=True))
model_zoo.register(name='transformers_gpt_double_heads',
model_fn=lambda: transformers.GPT2DoubleHeadsModel(config),
data_gen_fn=data_gen,
output_transform_fn=output_transform_fn,
model_attribute=ModelAttribute(has_control_flow=True))
model_zoo.register(name='transformers_gpt_for_token_classification',
model_fn=lambda: transformers.GPT2ForTokenClassification(config),
data_gen_fn=data_gen,
output_transform_fn=output_transform_fn,
model_attribute=ModelAttribute(has_control_flow=True))
model_zoo.register(name='transformers_gpt_for_sequence_classification',
model_fn=lambda: transformers.GPT2ForSequenceClassification(config),
data_gen_fn=data_gen,
output_transform_fn=output_transform_fn,
model_attribute=ModelAttribute(has_control_flow=True))