From e14e948e13d969a78736c4db34a3aa4b0de59533 Mon Sep 17 00:00:00 2001
From: Tim O'Donnell <timodonnell@gmail.com>
Date: Mon, 2 Sep 2019 11:01:58 -0400
Subject: [PATCH] add MHCFLURRY_DEFAULT_PREDICT_BATCH_SIZE environment variable
---
mhcflurry/class1_neural_network.py | 14 ++++++++++++--
1 file changed, 12 insertions(+), 2 deletions(-)
diff --git a/mhcflurry/class1_neural_network.py b/mhcflurry/class1_neural_network.py
index 0be671ac..1ebf7230 100644
--- a/mhcflurry/class1_neural_network.py
+++ b/mhcflurry/class1_neural_network.py
@@ -1,9 +1,10 @@
import time
import collections
-import logging
import json
import weakref
import itertools
+import os
+import logging
import numpy
import pandas
@@ -17,6 +18,15 @@ from .custom_loss import get_loss
from .data_dependent_weights_initialization import lsuv_init
+DEFAULT_PREDICT_BATCH_SIZE = 4096
+if os.environ.get("MHCFLURRY_DEFAULT_PREDICT_BATCH_SIZE"):
+ DEFAULT_PREDICT_BATCH_SIZE = os.environ[
+ "MHCFLURRY_DEFAULT_PREDICT_BATCH_SIZE"
+ ]
+ logging.info(
+ "Configured default predict batch size: %d" % DEFAULT_PREDICT_BATCH_SIZE)
+
+
class Class1NeuralNetwork(object):
"""
Low level class I predictor consisting of a single neural network.
@@ -1001,7 +1011,7 @@ class Class1NeuralNetwork(object):
self,
peptides,
allele_encoding=None,
- batch_size=4096,
+ batch_size=DEFAULT_PREDICT_BATCH_SIZE,
output_index=0):
"""
Predict affinities.
--
GitLab