diff --git a/mhcflurry/class1_neural_network.py b/mhcflurry/class1_neural_network.py index 0be671aca5634d988114722396aabcb6a533c40c..1ebf7230e6a2b146f6345c77c7bf81f9b61866a9 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.