Merge pull request #253 from CyrusBiotechnology/long_sequence_inference

Add long_sequence_inference option to command line

README.md

@@ -229,7 +229,8 @@ Using the most conservative settings, we were able to run inference on a
 offloading mode, ours is considerably faster; the same complex takes the more 
 efficent AlphaFold-Multimer more than double the time. Use the
 `long_sequence_inference` config option to enable all of these interventions
-at once.
+at once. The `run_pretrained_openfold.py` script can enable this config option with the 
+`--long_sequence_inference` command line option
 
 ### Training
 

run_pretrained_openfold.py

@@ -140,7 +140,7 @@ def main(args):
     # Create the output directory
     os.makedirs(args.output_dir, exist_ok=True)
 
-    config = model_config(args.config_preset)
+    config = model_config(args.config_preset, long_sequence_inference=args.long_sequence_inference)
     
     if(args.trace_model):
         if(not config.data.predict.fixed_size):
@@ -369,6 +369,10 @@ if __name__ == "__main__":
         help=""""Whether to output (100 - pLDDT) in the B-factor column instead
                  of the pLDDT itself"""
     )
+    parser.add_argument(
+        "--long_sequence_inference", action="store_true", default=False,
+        help="""enable options to reduce memory usage at the cost of speed, helps longer sequences fit into GPU memory, see the README for details"""
+    )
     add_data_args(parser)
     args = parser.parse_args()