github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-05 19:13:01 +00:00
Code Issues Packages Projects Releases Wiki Activity

[booster] add low level zero plugin (#3594)

Browse Source 
* [booster] add low level zero plugin

* [booster] fix gemini plugin test

* [booster] fix precision

* [booster] add low level zero plugin test

* [test] fix booster plugin test oom

* [test] fix booster plugin test oom

* [test] fix googlenet and inception output trans

* [test] fix diffuser clip vision model

* [test] fix torchaudio_wav2vec2_base

* [test] fix low level zero plugin test
This commit is contained in:
Hongxin Liu
2023-04-26 14:37:25 +08:00
committed by GitHub
parent b9a8dff7e5
commit 4b3240cb59
9 changed files with 476 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
tests/kit/model_zoo/torchvision/torchvision.py
View File

@@ -36,12 +36,12 @@ def swin_s():
# special output transform fn
google_net_output_transform_fn = lambda x: dict(output=x.logits) if isinstance(x, torchvision.models.GoogLeNetOutputs
) else dict(output=x)
google_net_output_transform_fn = lambda x: dict(output=sum(x)) if isinstance(x, torchvision.models.GoogLeNetOutputs
) else dict(output=x)
swin_s_output_output_transform_fn = lambda x: {f'output{idx}': val
for idx, val in enumerate(x)} if isinstance(x, tuple) else dict(output=x)
inception_v3_output_transform_fn = lambda x: dict(output=x.logits) if isinstance(x, torchvision.models.InceptionOutputs
) else dict(output=x)
inception_v3_output_transform_fn = lambda x: dict(output=sum(x)) if isinstance(x, torchvision.models.InceptionOutputs
) else dict(output=x)
model_zoo.register(name='torchvision_alexnet',
model_fn=tm.alexnet,