Skip to content

GitLab

Commit 71e72c48 authored by oahzxl's avatar oahzxl
Browse files

last version of benchmark

parent 55cb713f
Hide whitespace changes
Inline Side-by-side
Showing with 6 additions and 4 deletions
autochunk_benchmark.py
View file @ 71e72c48
...@@ -93,22 +93,24 @@ def _build_openfold(): ...@@ -93,22 +93,24 @@ def _build_openfold():
def benchmark_evoformer(): def benchmark_evoformer():
# init data and model # init data and model
msa_len = 300 msa_len = 256
pair_len = 800 pair_len = 2048
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 = 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) autochunk = _build_autochunk(evoformer_base().cuda(), max_memory, node, pair)
# build openfold # build openfold
chunk_size = 64
openfold = _build_openfold() openfold = _build_openfold()
# benchmark # benchmark
_benchmark_evoformer(model, node, pair, "base") _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") _benchmark_evoformer(autochunk, node, pair, "autochunk")
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment