Skip to content

Commit cdd06ee

Browse files
lvhoaalvhoaa
committed
update packages imports
1 parent e3bee4f commit cdd06ee

File tree

8 files changed

+45
-138
lines changed

8 files changed

+45
-138
lines changed

‎README.md‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ Why MegaFold?
2626

2727
## Usage
2828

29-
We include code for AlphaFold3 training (both baseline and end-to-end MegaFold integrations) and instructions to reproduce our paper results.
29+
We include code for AlphaFold3 training with end-to-end MegaFold integrations and instructions to reproduce our paper results.
3030

3131
### Install required dependencies
3232

@@ -104,7 +104,7 @@ rm input_cache
104104
### Run code
105105

106106
```
107-
AF3_OPTIMIZATIONS_MODE="megafold" python3 train.py --config configs/megafold_interactive.yaml --trainer_name initial_training
107+
python3 train.py --config configs/megafold_interactive.yaml --trainer_name initial_training
108108
```
109109

110110
Script to submit batch jobs is available in `scripts`. For example, you want to launch a job with `nodes=1` and `gpus=2`:

‎megafold/__init__.py‎

Lines changed: 28 additions & 87 deletions
Original file line numberDiff line numberDiff line change
@@ -30,94 +30,35 @@
3030
register_input_transform,
3131
)
3232
import os
33-
if os.environ.get("AF3_OPTIMIZATIONS_MODE") == "baseline":
34-
from megafold.baseline import (
35-
AdaptiveLayerNorm,
36-
AttentionPairBias,
37-
CentreRandomAugmentation,
38-
ComputeModelSelectionScore,
39-
ComputeRankingScore,
40-
ConditionWrapper,
41-
ConfidenceHead,
42-
ConfidenceHeadLogits,
43-
DiffusionModule,
44-
DiffusionTransformer,
45-
DistogramHead,
46-
ElucidatedAtomDiffusion,
47-
InputFeatureEmbedder,
48-
MSAModule,
49-
MSAPairWeightedAveraging,
50-
MultiChainPermutationAlignment,
51-
MegaFold,
52-
MegaFoldWithHubMixin,
53-
OuterProductMean,
54-
PairformerStack,
55-
PreLayerNorm,
56-
RelativePositionEncoding,
57-
TemplateEmbedder,
58-
Transition,
59-
TriangleAttention,
60-
TriangleMultiplication,
61-
)
62-
elif os.environ.get("AF3_OPTIMIZATIONS_MODE") == "layernormlinear":
63-
from megafold.megafold_layernormlinear import (
64-
AdaptiveLayerNorm,
65-
AttentionPairBias,
66-
CentreRandomAugmentation,
67-
ComputeModelSelectionScore,
68-
ComputeRankingScore,
69-
ConditionWrapper,
70-
ConfidenceHead,
71-
ConfidenceHeadLogits,
72-
DiffusionModule,
73-
DiffusionTransformer,
74-
DistogramHead,
75-
ElucidatedAtomDiffusion,
76-
InputFeatureEmbedder,
77-
MSAModule,
78-
MSAPairWeightedAveraging,
79-
MultiChainPermutationAlignment,
80-
MegaFold,
81-
MegaFoldWithHubMixin,
82-
OuterProductMean,
83-
PairformerStack,
84-
PreLayerNorm,
85-
RelativePositionEncoding,
86-
TemplateEmbedder,
87-
Transition,
88-
TriangleAttention,
89-
TriangleMultiplication,
90-
)
91-
else:
92-
from megafold.model.megafold import (
93-
AdaptiveLayerNorm,
94-
AttentionPairBias,
95-
CentreRandomAugmentation,
96-
ComputeModelSelectionScore,
97-
ComputeRankingScore,
98-
ConditionWrapper,
99-
ConfidenceHead,
100-
ConfidenceHeadLogits,
101-
DiffusionModule,
102-
DiffusionTransformer,
103-
DistogramHead,
104-
ElucidatedAtomDiffusion,
105-
InputFeatureEmbedder,
106-
MSAModule,
107-
MSAPairWeightedAveraging,
108-
MultiChainPermutationAlignment,
109-
MegaFold,
110-
MegaFoldWithHubMixin,
111-
OuterProductMean,
112-
PairformerStack,
113-
PreLayerNorm,
114-
RelativePositionEncoding,
115-
TemplateEmbedder,
116-
Transition,
117-
TriangleAttention,
118-
TriangleMultiplication,
119-
)
12033

34+
from megafold.model.megafold import (
35+
AdaptiveLayerNorm,
36+
AttentionPairBias,
37+
CentreRandomAugmentation,
38+
ComputeModelSelectionScore,
39+
ComputeRankingScore,
40+
ConditionWrapper,
41+
ConfidenceHead,
42+
ConfidenceHeadLogits,
43+
DiffusionModule,
44+
DiffusionTransformer,
45+
DistogramHead,
46+
ElucidatedAtomDiffusion,
47+
InputFeatureEmbedder,
48+
MSAModule,
49+
MSAPairWeightedAveraging,
50+
MultiChainPermutationAlignment,
51+
MegaFold,
52+
MegaFoldWithHubMixin,
53+
OuterProductMean,
54+
PairformerStack,
55+
PreLayerNorm,
56+
RelativePositionEncoding,
57+
TemplateEmbedder,
58+
Transition,
59+
TriangleAttention,
60+
TriangleMultiplication,
61+
)
12162
from megafold.trainer import DataLoader, Trainer
12263
from megafold.utils.model_utils import (
12364
ComputeAlignmentError,

‎megafold/app.py‎

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -6,13 +6,8 @@
66
from Bio.PDB import PDBIO
77

88
from megafold.inputs import MegaFoldInput
9-
import os
10-
if os.environ.get("AF3_OPTIMIZATIONS_MODE") == "baseline":
11-
from megafold.baseline import MegaFold
12-
elif os.environ.get("AF3_OPTIMIZATIONS_MODE") == "layernormlinear":
13-
from megafold.megafold_layernormlinear import MegaFold
14-
else:
15-
from megafold.model.megafold import MegaFold
9+
import os
10+
from megafold.model.megafold import MegaFold
1611

1712
# constants
1813
model = None

‎megafold/cli.py‎

Lines changed: 7 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -17,27 +17,13 @@
1717
MegaFoldInput,
1818
megafold_input_to_pdb_input,
1919
)
20-
if os.environ.get("AF3_OPTIMIZATIONS_MODE") == "baseline":
21-
from megafold.baseline import (
22-
ComputeConfidenceScore,
23-
ComputeRankingScore,
24-
ConfidenceHeadLogits,
25-
MegaFold,
26-
)
27-
elif os.environ.get("AF3_OPTIMIZATIONS_MODE") == "layernormlinear":
28-
from megafold.megafold_layernormlinear import (
29-
ComputeConfidenceScore,
30-
ComputeRankingScore,
31-
ConfidenceHeadLogits,
32-
MegaFold,
33-
)
34-
else:
35-
from megafold.model.megafold import (
36-
ComputeConfidenceScore,
37-
ComputeRankingScore,
38-
ConfidenceHeadLogits,
39-
MegaFold,
40-
)
20+
21+
from megafold.model.megafold import (
22+
ComputeConfidenceScore,
23+
ComputeRankingScore,
24+
ConfidenceHeadLogits,
25+
MegaFold,
26+
)
4127
from megafold.tensor_typing import typecheck
4228
from megafold.utils.data_utils import decrement_all_by_n
4329
from megafold.utils.model_utils import batch_repeat_interleave, lens_to_mask, not_exists

‎megafold/configs.py‎

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -21,12 +21,7 @@
2121
)
2222
from megafold.nlm import NLMEmbedding
2323
import os
24-
if os.environ.get("AF3_OPTIMIZATIONS_MODE") == "baseline":
25-
from megafold.baseline import MegaFold
26-
elif os.environ.get("AF3_OPTIMIZATIONS_MODE") == "layernormlinear":
27-
from megafold.megafold_layernormlinear import MegaFold
28-
else:
29-
from megafold.model.megafold import MegaFold
24+
from megafold.model.megafold import MegaFold
3025
from megafold.plm import PLMEmbedding
3126
from megafold.tensor_typing import typecheck
3227
from megafold.trainer import Dataset, Fabric, LRScheduler, Optimizer, Trainer

‎megafold/inputs.py‎

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -119,13 +119,8 @@
119119

120120
PDB_INPUT_CHAIN_SEQS = List[List[Tuple[str, str, PDB_INPUT_RESIDUE_MOLECULE_TYPE]]]
121121

122-
if os.environ.get("AF3_OPTIMIZATIONS_MODE") == "baseline":
123-
# baseline does not have these logic related to cache&timeout -> set timeout to 300/100 here
124-
PDB_INPUT_TO_MOLECULE_INPUT_MAX_SECONDS_PER_INPUT = 300
125-
MOLECULE_TO_ATOM_INPUT_MAX_SECONDS_PER_INPUT = 100
126-
else:
127-
PDB_INPUT_TO_MOLECULE_INPUT_MAX_SECONDS_PER_INPUT = 30
128-
MOLECULE_TO_ATOM_INPUT_MAX_SECONDS_PER_INPUT = 10
122+
PDB_INPUT_TO_MOLECULE_INPUT_MAX_SECONDS_PER_INPUT = 30
123+
MOLECULE_TO_ATOM_INPUT_MAX_SECONDS_PER_INPUT = 10
129124

130125
IS_MOLECULE_TYPES = 5
131126
IS_PROTEIN_INDEX = 0

‎megafold/trainer.py‎

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -37,12 +37,7 @@
3737
collate_inputs_to_batched_atom_input,
3838
compose_calls,
3939
)
40-
if os.environ.get("AF3_OPTIMIZATIONS_MODE") == "baseline":
41-
from megafold.baseline import ComputeConfidenceScore, ComputeModelSelectionScore, Sample
42-
elif os.environ.get("AF3_OPTIMIZATIONS_MODE") == "layernormlinear":
43-
from megafold.megafold_layernormlinear import ComputeConfidenceScore, ComputeModelSelectionScore, Sample
44-
else:
45-
from megafold.model.megafold import ComputeConfidenceScore, ComputeModelSelectionScore, Sample
40+
from megafold.model.megafold import ComputeConfidenceScore, ComputeModelSelectionScore, Sample
4641

4742
from megafold.tensor_typing import package_available, should_typecheck, typecheck
4843
from megafold.utils.model_utils import at_most_one_of, divisible_by

‎scripts/megafold.sh‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,8 +27,8 @@ export NUM_NODES NUM_GPUS GLOBAL_BATCH_SIZE NUM_WORKERS MULTIPROCESSING_CONTEXT
2727

2828
envsubst < configs/megafold.yaml > configs/megafold_rendered.yaml
2929
if [ "$GLOBAL_BATCH_SIZE" -gt 1 ]; then
30-
AF3_OPTIMIZATIONS_MODE="megafold" srun python3 train.py --config configs/megafold_rendered.yaml --trainer_name initial_training
30+
srun python3 train.py --config configs/megafold_rendered.yaml --trainer_name initial_training
3131
else
32-
AF3_OPTIMIZATIONS_MODE="megafold" python3 train.py --config configs/megafold_rendered.yaml --trainer_name initial_training
32+
python3 train.py --config configs/megafold_rendered.yaml --trainer_name initial_training
3333
fi
3434

0 commit comments

Comments
 (0)