github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

set benchmark

pull/2364/head
oahzxl 2023-01-09 14:20:41 +08:00
parent 4d223e18a2
commit cb68ee864a
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
tests/test_autochunk/benchmark_autochunk.py
View File

@ -98,14 +98,14 @@ def _build_openfold():
def benchmark_evoformer():
# init data and model
msa_len = 256
pair_len = 256
pair_len = 512
node = torch.randn(1, msa_len, pair_len, 256).cuda()
pair = torch.randn(1, pair_len, pair_len, 128).cuda()
model = evoformer_base().cuda()
# build autochunk model
max_memory = 1000 # MB fit memory mode
# max_memory = None # min memory mode
# max_memory = 1000 # MB, fit memory mode
max_memory = None # min memory mode
autochunk = _build_autochunk(evoformer_base().cuda(), max_memory, node, pair)
# build openfold