From 71e72c48907195096ef02be73e1c5b0feea2653d Mon Sep 17 00:00:00 2001 From: oahzxl Date: Thu, 5 Jan 2023 17:54:25 +0800 Subject: [PATCH] last version of benchmark --- autochunk_benchmark.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/autochunk_benchmark.py b/autochunk_benchmark.py index 3b48d7e46..c938485ef 100644 --- a/autochunk_benchmark.py +++ b/autochunk_benchmark.py @@ -93,22 +93,24 @@ def _build_openfold(): def benchmark_evoformer(): # init data and model - msa_len = 300 - pair_len = 800 + msa_len = 256 + pair_len = 2048 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 = 3000 # MB + max_memory = 10000 # MB fit memory mode + # max_memory = None # min memory mode autochunk = _build_autochunk(evoformer_base().cuda(), max_memory, node, pair) # build openfold + chunk_size = 64 openfold = _build_openfold() # benchmark _benchmark_evoformer(model, node, pair, "base") - _benchmark_evoformer(openfold, node, pair, "openfold", chunk_size=4) + _benchmark_evoformer(openfold, node, pair, "openfold", chunk_size=chunk_size) _benchmark_evoformer(autochunk, node, pair, "autochunk")