From efb1c64c30cf2ee35dad03bfd3829f014d204a8d Mon Sep 17 00:00:00 2001 From: oahzxl Date: Fri, 6 Jan 2023 11:39:26 +0800 Subject: [PATCH] restruct dir --- .../autochunk}/chunk_codegen.py | 0 .../test_autochunk/autochunk_benchmark.py | 14 +++++++------- .../test_autochunk/autochunk_test.py | 4 ++-- .../test_autochunk}/evoformer/evoformer.py | 0 .../test_autochunk}/evoformer/initializer.py | 0 .../test_autochunk}/evoformer/kernel.py | 0 .../test_autochunk}/evoformer/msa.py | 0 .../test_autochunk}/evoformer/ops.py | 0 .../test_autochunk}/evoformer/triangle.py | 0 .../test_autochunk}/openfold/checkpointing.py | 0 .../test_autochunk}/openfold/dropout.py | 0 .../test_autochunk}/openfold/evoformer.py | 18 +++++++++--------- .../test_autochunk}/openfold/msa.py | 6 +++--- .../openfold/outer_product_mean.py | 4 ++-- .../openfold/pair_transition.py | 4 ++-- .../test_autochunk}/openfold/primitives.py | 4 ++-- .../test_autochunk}/openfold/tensor_utils.py | 0 .../openfold/triangular_attention.py | 4 ++-- .../triangular_multiplicative_update.py | 4 ++-- 19 files changed, 31 insertions(+), 31 deletions(-) rename {autochunk => colossalai/autochunk}/chunk_codegen.py (100%) rename autochunk_benchmark.py => tests/test_autochunk/autochunk_benchmark.py (89%) rename autochunk_test.py => tests/test_autochunk/autochunk_test.py (96%) rename {autochunk => tests/test_autochunk}/evoformer/evoformer.py (100%) rename {autochunk => tests/test_autochunk}/evoformer/initializer.py (100%) rename {autochunk => tests/test_autochunk}/evoformer/kernel.py (100%) rename {autochunk => tests/test_autochunk}/evoformer/msa.py (100%) rename {autochunk => tests/test_autochunk}/evoformer/ops.py (100%) rename {autochunk => tests/test_autochunk}/evoformer/triangle.py (100%) rename {autochunk => tests/test_autochunk}/openfold/checkpointing.py (100%) rename {autochunk => tests/test_autochunk}/openfold/dropout.py (100%) rename {autochunk => tests/test_autochunk}/openfold/evoformer.py (96%) rename {autochunk => tests/test_autochunk}/openfold/msa.py (98%) rename {autochunk => tests/test_autochunk}/openfold/outer_product_mean.py (97%) rename {autochunk => tests/test_autochunk}/openfold/pair_transition.py (96%) rename {autochunk => tests/test_autochunk}/openfold/primitives.py (99%) rename {autochunk => tests/test_autochunk}/openfold/tensor_utils.py (100%) rename {autochunk => tests/test_autochunk}/openfold/triangular_attention.py (97%) rename {autochunk => tests/test_autochunk}/openfold/triangular_multiplicative_update.py (97%) diff --git a/autochunk/chunk_codegen.py b/colossalai/autochunk/chunk_codegen.py similarity index 100% rename from autochunk/chunk_codegen.py rename to colossalai/autochunk/chunk_codegen.py diff --git a/autochunk_benchmark.py b/tests/test_autochunk/autochunk_benchmark.py similarity index 89% rename from autochunk_benchmark.py rename to tests/test_autochunk/autochunk_benchmark.py index c34b5217e..8df6d9ff4 100644 --- a/autochunk_benchmark.py +++ b/tests/test_autochunk/autochunk_benchmark.py @@ -3,13 +3,13 @@ import time import torch import torch.fx -from autochunk.chunk_codegen import ChunkCodeGen +from colossalai.autochunk.chunk_codegen import ChunkCodeGen from colossalai.fx import ColoTracer from colossalai.fx.graph_module import ColoGraphModule from colossalai.fx.passes.meta_info_prop import MetaInfoProp from colossalai.fx.profiler import MetaTensor -from autochunk.evoformer.evoformer import evoformer_base -from autochunk.openfold.evoformer import EvoformerBlock +from tests.test_autochunk.evoformer.evoformer import evoformer_base +from tests.test_autochunk.openfold.evoformer import EvoformerBlock def _benchmark_evoformer(model: torch.nn.Module, node, pair, title, chunk_size=None): @@ -94,7 +94,7 @@ def _build_openfold(): def benchmark_evoformer(): # init data and model msa_len = 256 - pair_len = 1024 + pair_len = 256 node = torch.randn(1, msa_len, pair_len, 256).cuda() pair = torch.randn(1, pair_len, pair_len, 128).cuda() model = evoformer_base().cuda() @@ -106,11 +106,11 @@ def benchmark_evoformer(): # build openfold chunk_size = 64 - # openfold = _build_openfold() + openfold = _build_openfold() # benchmark - # _benchmark_evoformer(model, node, pair, "base") - # _benchmark_evoformer(openfold, node, pair, "openfold", chunk_size=chunk_size) + _benchmark_evoformer(model, node, pair, "base") + _benchmark_evoformer(openfold, node, pair, "openfold", chunk_size=chunk_size) _benchmark_evoformer(autochunk, node, pair, "autochunk") diff --git a/autochunk_test.py b/tests/test_autochunk/autochunk_test.py similarity index 96% rename from autochunk_test.py rename to tests/test_autochunk/autochunk_test.py index 63f393531..5e9aaca15 100644 --- a/autochunk_test.py +++ b/tests/test_autochunk/autochunk_test.py @@ -12,8 +12,8 @@ from colossalai.core import global_context as gpc from colossalai.fx.graph_module import ColoGraphModule from colossalai.fx.passes.meta_info_prop import MetaInfoProp, TensorMetadata from colossalai.fx.profiler import MetaTensor -from autochunk.evoformer.evoformer import evoformer_base -from autochunk.chunk_codegen import ChunkCodeGen +from tests.test_autochunk.evoformer.evoformer import evoformer_base +from ...colossalai.autochunk.chunk_codegen import ChunkCodeGen with_codegen = True diff --git a/autochunk/evoformer/evoformer.py b/tests/test_autochunk/evoformer/evoformer.py similarity index 100% rename from autochunk/evoformer/evoformer.py rename to tests/test_autochunk/evoformer/evoformer.py diff --git a/autochunk/evoformer/initializer.py b/tests/test_autochunk/evoformer/initializer.py similarity index 100% rename from autochunk/evoformer/initializer.py rename to tests/test_autochunk/evoformer/initializer.py diff --git a/autochunk/evoformer/kernel.py b/tests/test_autochunk/evoformer/kernel.py similarity index 100% rename from autochunk/evoformer/kernel.py rename to tests/test_autochunk/evoformer/kernel.py diff --git a/autochunk/evoformer/msa.py b/tests/test_autochunk/evoformer/msa.py similarity index 100% rename from autochunk/evoformer/msa.py rename to tests/test_autochunk/evoformer/msa.py diff --git a/autochunk/evoformer/ops.py b/tests/test_autochunk/evoformer/ops.py similarity index 100% rename from autochunk/evoformer/ops.py rename to tests/test_autochunk/evoformer/ops.py diff --git a/autochunk/evoformer/triangle.py b/tests/test_autochunk/evoformer/triangle.py similarity index 100% rename from autochunk/evoformer/triangle.py rename to tests/test_autochunk/evoformer/triangle.py diff --git a/autochunk/openfold/checkpointing.py b/tests/test_autochunk/openfold/checkpointing.py similarity index 100% rename from autochunk/openfold/checkpointing.py rename to tests/test_autochunk/openfold/checkpointing.py diff --git a/autochunk/openfold/dropout.py b/tests/test_autochunk/openfold/dropout.py similarity index 100% rename from autochunk/openfold/dropout.py rename to tests/test_autochunk/openfold/dropout.py diff --git a/autochunk/openfold/evoformer.py b/tests/test_autochunk/openfold/evoformer.py similarity index 96% rename from autochunk/openfold/evoformer.py rename to tests/test_autochunk/openfold/evoformer.py index ffd4c9829..b53ec1aa5 100644 --- a/autochunk/openfold/evoformer.py +++ b/tests/test_autochunk/openfold/evoformer.py @@ -19,25 +19,25 @@ import torch.nn as nn from typing import Tuple, Optional from functools import partial -from openfold.primitives import Linear, LayerNorm -from openfold.dropout import DropoutRowwise, DropoutColumnwise -from openfold.msa import ( +from .primitives import Linear, LayerNorm +from .dropout import DropoutRowwise, DropoutColumnwise +from .msa import ( MSARowAttentionWithPairBias, MSAColumnAttention, MSAColumnGlobalAttention, ) -from openfold.outer_product_mean import OuterProductMean -from openfold.pair_transition import PairTransition -from openfold.triangular_attention import ( +from .outer_product_mean import OuterProductMean +from .pair_transition import PairTransition +from .triangular_attention import ( TriangleAttentionStartingNode, TriangleAttentionEndingNode, ) -from openfold.triangular_multiplicative_update import ( +from .triangular_multiplicative_update import ( TriangleMultiplicationOutgoing, TriangleMultiplicationIncoming, ) -from openfold.checkpointing import checkpoint_blocks, get_checkpoint_fn -from openfold.tensor_utils import chunk_layer +from .checkpointing import checkpoint_blocks, get_checkpoint_fn +from .tensor_utils import chunk_layer class MSATransition(nn.Module): diff --git a/autochunk/openfold/msa.py b/tests/test_autochunk/openfold/msa.py similarity index 98% rename from autochunk/openfold/msa.py rename to tests/test_autochunk/openfold/msa.py index 00b822e7f..7c137286f 100644 --- a/autochunk/openfold/msa.py +++ b/tests/test_autochunk/openfold/msa.py @@ -18,15 +18,15 @@ import torch import torch.nn as nn from typing import Optional, List, Tuple -from openfold.primitives import ( +from .primitives import ( Linear, LayerNorm, Attention, GlobalAttention, _attention_chunked_trainable, ) -from openfold.checkpointing import get_checkpoint_fn -from openfold.tensor_utils import ( +from .checkpointing import get_checkpoint_fn +from .tensor_utils import ( chunk_layer, permute_final_dims, flatten_final_dims, diff --git a/autochunk/openfold/outer_product_mean.py b/tests/test_autochunk/openfold/outer_product_mean.py similarity index 97% rename from autochunk/openfold/outer_product_mean.py rename to tests/test_autochunk/openfold/outer_product_mean.py index 43d853833..daadf1c27 100644 --- a/autochunk/openfold/outer_product_mean.py +++ b/tests/test_autochunk/openfold/outer_product_mean.py @@ -19,8 +19,8 @@ from typing import Optional import torch import torch.nn as nn -from openfold.primitives import Linear -from openfold.tensor_utils import chunk_layer +from .primitives import Linear +from .tensor_utils import chunk_layer class OuterProductMean(nn.Module): diff --git a/autochunk/openfold/pair_transition.py b/tests/test_autochunk/openfold/pair_transition.py similarity index 96% rename from autochunk/openfold/pair_transition.py rename to tests/test_autochunk/openfold/pair_transition.py index de7630641..7d09914dc 100644 --- a/autochunk/openfold/pair_transition.py +++ b/tests/test_autochunk/openfold/pair_transition.py @@ -17,8 +17,8 @@ from typing import Optional import torch import torch.nn as nn -from openfold.primitives import Linear, LayerNorm -from openfold.tensor_utils import chunk_layer +from .primitives import Linear, LayerNorm +from .tensor_utils import chunk_layer class PairTransition(nn.Module): diff --git a/autochunk/openfold/primitives.py b/tests/test_autochunk/openfold/primitives.py similarity index 99% rename from autochunk/openfold/primitives.py rename to tests/test_autochunk/openfold/primitives.py index bbc156f21..32a9d487c 100644 --- a/autochunk/openfold/primitives.py +++ b/tests/test_autochunk/openfold/primitives.py @@ -21,8 +21,8 @@ import numpy as np import torch import torch.nn as nn -from openfold.checkpointing import get_checkpoint_fn -from openfold.tensor_utils import ( +from .checkpointing import get_checkpoint_fn +from .tensor_utils import ( permute_final_dims, flatten_final_dims, _chunk_slice, diff --git a/autochunk/openfold/tensor_utils.py b/tests/test_autochunk/openfold/tensor_utils.py similarity index 100% rename from autochunk/openfold/tensor_utils.py rename to tests/test_autochunk/openfold/tensor_utils.py diff --git a/autochunk/openfold/triangular_attention.py b/tests/test_autochunk/openfold/triangular_attention.py similarity index 97% rename from autochunk/openfold/triangular_attention.py rename to tests/test_autochunk/openfold/triangular_attention.py index 6d3e37f4c..12d09c502 100644 --- a/autochunk/openfold/triangular_attention.py +++ b/tests/test_autochunk/openfold/triangular_attention.py @@ -20,8 +20,8 @@ from typing import Optional, List import torch import torch.nn as nn -from openfold.primitives import Linear, LayerNorm, Attention -from openfold.tensor_utils import ( +from .primitives import Linear, LayerNorm, Attention +from .tensor_utils import ( chunk_layer, permute_final_dims, flatten_final_dims, diff --git a/autochunk/openfold/triangular_multiplicative_update.py b/tests/test_autochunk/openfold/triangular_multiplicative_update.py similarity index 97% rename from autochunk/openfold/triangular_multiplicative_update.py rename to tests/test_autochunk/openfold/triangular_multiplicative_update.py index 2406e2bac..29f7062c3 100644 --- a/autochunk/openfold/triangular_multiplicative_update.py +++ b/tests/test_autochunk/openfold/triangular_multiplicative_update.py @@ -19,8 +19,8 @@ from typing import Optional import torch import torch.nn as nn -from openfold.primitives import Linear, LayerNorm -from openfold.tensor_utils import permute_final_dims +from .primitives import Linear, LayerNorm +from .tensor_utils import permute_final_dims class TriangleMultiplicativeUpdate(nn.Module):