Add better GPU support

PredictionServer/README

@@ -1,7 +1,8 @@
 Solubility / Usability prediction server
 
-For installing required packages use:
-pip install -r requirements.txt
+For installing required packages use (for inference on CPU or GPU, respectively):
+pip install -r requirements_cpu.txt
+pip install -r requirements_gpu.txt
 
 All models go into models/
 Model name format: {PREDICTION_TYPE}_{MODEL_TYPE}_{Fold number 0-4}_quantized.onnx

PredictionServer/predict.py

@@ -49,11 +49,12 @@ def run_model_distilled(embed_dataloader, args, prediction_type, test_df):
     opts.intra_op_num_threads = args.NUM_THREADS
     opts.inter_op_num_threads = args.NUM_THREADS
     opts.execution_mode = onnxruntime.ExecutionMode.ORT_SEQUENTIAL
-
+    providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] if torch.cuda.is_available() else ["CPUExecutionProvider"]
+
     # Adjust session options
     model_paths = [os.path.join(args.MODELS_PATH,
       f"{prediction_type}_ESM1b_distilled_quantized.onnx")]
-    ort_sessions = [onnxruntime.InferenceSession(mp, sess_options=opts) for mp in model_paths]
+    ort_sessions = [onnxruntime.InferenceSession(mp, sess_options=opts, providers=providers) for mp in model_paths]
 
     embed_dict = {}
     inputs_names = ort_sessions[0].get_inputs()

PredictionServer/requirements.txt → PredictionServer/requirements_cpu.txt

@@ -1,6 +1,5 @@
 torch>=1.6
 onnxruntime>=1.7.0
-onnxruntime-gpu
 numpy
 pandas
 fair-esm

PredictionServer/requirements_gpu.txt

@@ -0,0 +1,5 @@
+torch>=1.6
+onnxruntime-gpu>=1.7.0
+numpy
+pandas
+fair-esm