github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-06-22 21:49:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

set benchmark

Browse Source
This commit is contained in:
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(): def benchmark_evoformer():
# init data and model # init data and model
msa_len = 256 msa_len = 256
pair_len = 256 pair_len = 512
node = torch.randn(1, msa_len, pair_len, 256).cuda() node = torch.randn(1, msa_len, pair_len, 256).cuda()
pair = torch.randn(1, pair_len, pair_len, 128).cuda() pair = torch.randn(1, pair_len, pair_len, 128).cuda()
model = evoformer_base().cuda() model = evoformer_base().cuda()
# build autochunk model # build autochunk model
max_memory = 1000 # MB fit memory mode # max_memory = 1000 # MB, fit memory mode
# max_memory = None # min memory mode max_memory = None # min memory mode
autochunk = _build_autochunk(evoformer_base().cuda(), max_memory, node, pair) autochunk = _build_autochunk(evoformer_base().cuda(), max_memory, node, pair)
# build openfold # build openfold