implement merging optimization

mhcflurry/class1_affinity_predictor.py

@@ -5,7 +5,7 @@ import logging
 import time
 import warnings
 from os.path import join, exists, abspath
-from os import mkdir
+from os import mkdir, environ
 from socket import gethostname
 from getpass import getuser
 from functools import partial
@@ -31,6 +31,9 @@ from .allele_encoding import AlleleEncoding
 # See ensemble_centrality.py for other options.
 DEFAULT_CENTRALITY_MEASURE = "mean"
 
+# Any value > 0 will result in attempting to optimize models after loading.
+OPTIMIZATION_LEVEL = int(environ.get("MHCFLURRY_OPTIMIZATION_LEVEL", 0))
+
 
 class Class1AffinityPredictor(object):
     """
@@ -98,6 +101,7 @@ class Class1AffinityPredictor(object):
         self.metadata_dataframes = (
             dict(metadata_dataframes) if metadata_dataframes else {})
         self._cache = {}
+        self.optimization_info = {}
 
         assert isinstance( self.allele_to_allele_specific_models, dict)
         assert isinstance(self.class1_pan_allele_models, list)
@@ -483,8 +487,48 @@ class Class1AffinityPredictor(object):
             manifest_df=manifest_df,
             allele_to_percent_rank_transform=allele_to_percent_rank_transform,
         )
+        if OPTIMIZATION_LEVEL >= 1:
+            logging.info("Optimizing models")
+            optimized = result.optimize()
+            logging.info(
+                "Optimization " + ("succeeded" if optimized else "failed"))
         return result
 
+    def optimize(self):
+        """
+        EXPERIMENTAL: Optimize the predictor for faster predictions.
+
+        Currently the only optimization implemented is to merge multiple pan-
+        allele predictors at the tensorflow level.
+
+        The optimization is performed in-place, mutating the instance.
+
+        Returns
+        ----------
+        bool
+            Whether optimization was performed
+
+        """
+        num_class1_pan_allele_models = len(self.class1_pan_allele_models)
+        if num_class1_pan_allele_models > 1:
+            try:
+                self.class1_pan_allele_models = [
+                    Class1NeuralNetwork.merge(
+                        self.class1_pan_allele_models,
+                        merge_method="concatenate")
+                ]
+            except NotImplementedError as e:
+                logging.warning("Optimization failed: %s" % str(e))
+                return False
+            self._manifest_df = None
+            self.clear_cache()
+            self.optimization_info["pan_models_merged"] = True
+            self.optimization_info["num_pan_models_merged"] = (
+                num_class1_pan_allele_models)
+        else:
+            return False
+        return True
+
     @staticmethod
     def model_name(allele, num):
         """
@@ -987,7 +1031,10 @@ class Class1AffinityPredictor(object):
             df["supported_peptide_length"] = True
             all_peptide_lengths_supported = True
 
-        num_pan_models = len(self.class1_pan_allele_models)
+        num_pan_models = (
+            len(self.class1_pan_allele_models)
+            if not self.optimization_info.get("pan_models_merged", False)
+            else self.optimization_info["num_pan_models_merged"])
         max_single_allele_models = max(
             len(self.allele_to_allele_specific_models.get(allele, []))
             for allele in unique_alleles
@@ -1015,40 +1062,46 @@ class Class1AffinityPredictor(object):
                     raise ValueError(msg)
             mask = df.supported_peptide_length & (
                 ~df.normalized_allele.isin(unsupported_alleles))
+
+            row_slice = None
             if mask is None or mask.all():
-                # Common case optimization
-                allele_encoding = AlleleEncoding(
+                row_slice = slice(None, None, None)  # all rows
+                masked_allele_encoding = AlleleEncoding(
                     df.normalized_allele,
                     borrow_from=master_allele_encoding)
+                masked_peptides = peptides
+            elif mask.sum() > 0:
+                row_slice = mask
+                masked_allele_encoding = AlleleEncoding(
+                    df.loc[mask].normalized_allele,
+                    borrow_from=master_allele_encoding)
+                masked_peptides = peptides.sequences[mask]
 
+            if row_slice is not None:
                 # The following line is a performance optimization that may be
                 # revisited. It causes the neural network to set to include
                 # only the alleles actually being predicted for. This makes
                 # the network much smaller. However, subsequent calls to
                 # predict will need to reset these weights, so there is a
                 # tradeoff.
-                allele_encoding = allele_encoding.compact()
-
-                for (i, model) in enumerate(self.class1_pan_allele_models):
-                    predictions_array[:, i] = (
-                        model.predict(
-                            peptides,
-                            allele_encoding=allele_encoding,
-                            **model_kwargs))
-            elif mask.sum() > 0:
-                masked_allele_encoding = AlleleEncoding(
-                    df.loc[mask].normalized_allele,
-                    borrow_from=master_allele_encoding)
-
-                # See above performance note.
                 masked_allele_encoding = masked_allele_encoding.compact()
 
-                masked_peptides = peptides.sequences[mask]
-                for (i, model) in enumerate(self.class1_pan_allele_models):
-                    predictions_array[mask, i] = model.predict(
+                if self.optimization_info.get("pan_models_merged"):
+                    # Multiple pan-allele models have been merged into one
+                    # at the tensorflow level.
+                    assert len(self.class1_pan_allele_models) == 1
+                    predictions = self.class1_pan_allele_models[0].predict(
                         masked_peptides,
                         allele_encoding=masked_allele_encoding,
+                        output_index=None,
                         **model_kwargs)
+                    predictions_array[row_slice, :num_pan_models] = predictions
+                else:
+                    for (i, model) in enumerate(self.class1_pan_allele_models):
+                        predictions_array[row_slice, i] = model.predict(
+                            masked_peptides,
+                            allele_encoding=masked_allele_encoding,
+                            **model_kwargs)
 
         if self.allele_to_allele_specific_models:
             unsupported_alleles = [

mhcflurry/class1_neural_network.py

@@ -5,6 +5,7 @@ import weakref
 import itertools
 import os
 import logging
+import pickle
 
 import numpy
 import pandas
@@ -1063,6 +1064,119 @@ class Class1NeuralNetwork(object):
             self.prediction_cache[peptides] = result
         return result
 
+    @classmethod
+    def merge(cls, models, merge_method="average"):
+        """
+        Merge multiple models at the tensorflow (or other backend) level.
+
+        Only certain neural network architectures support merging. Others will
+        throw NotImplementedError.
+
+        Parameters
+        ----------
+        models : list of Class1NeuralNetwork
+            instances to merge
+        merge_method : string, one of "average", "sum", or "concatenate"
+            How to merge the predictions of the different models
+
+        Returns
+        -------
+        Class1NeuralNetwork
+            The merged neural network
+
+        """
+        import keras
+        from keras.layers import Input
+        from keras.models import Model
+
+        if len(models) == 1:
+            return models[0]
+
+        # Copy models since we are going to mutate their underlying networks
+        models = [
+            pickle.loads(pickle.dumps(model, protocol=pickle.HIGHEST_PROTOCOL))
+            for model in models
+        ]
+        assert len(models) > 1
+
+        result = Class1NeuralNetwork(**dict(models[0].hyperparameters))
+
+        # Remove hyperparameters that are not shared by all models.
+        for model in models:
+            for (key, value) in model.hyperparameters.items():
+                if result.hyperparameters.get(key, value) != value:
+                    del result.hyperparameters[key]
+
+        assert result._network is None
+
+        networks = [
+            model.network() for model in models
+        ]
+
+        layer_names = [
+            [layer.name for layer in network.layers]
+            for network in networks
+        ]
+
+        pan_allele_layer_names1 = [
+            'allele', 'peptide', 'allele_representation', 'flattened_0',
+            'allele_flat', 'allele_peptide_merged', 'dense_0', 'dropout_0',
+            'dense_1', 'dropout_1', 'output',
+        ]
+
+        if all(names == pan_allele_layer_names1 for names in layer_names):
+            # Merging an ensemble of pan-allele architectures
+            network = networks[0]
+            peptide_input = Input(
+                shape=tuple(int(x) for x in network.inputs[0].shape[1:]),
+                dtype='float32',
+                name='peptide')
+            allele_input = Input(
+                shape=(1,),
+                dtype='float32',
+                name='allele')
+
+            allele_embedding = network.get_layer(
+                "allele_representation")(allele_input)
+            peptide_flat = network.get_layer("flattened_0")(peptide_input)
+            allele_flat = network.get_layer("allele_flat")(allele_embedding)
+            allele_peptide_merged = network.get_layer("allele_peptide_merged")(
+                [peptide_flat, allele_flat])
+
+            sub_networks = []
+            for (i, network) in enumerate(networks):
+                layers = network.layers[
+                    pan_allele_layer_names1.index("allele_peptide_merged") + 1:
+                ]
+                node = allele_peptide_merged
+                for layer in layers:
+                    layer.name += "_%d" % i
+                    node = layer(node)
+                sub_networks.append(node)
+
+            if merge_method == 'average':
+                output = keras.layers.average(sub_networks)
+            elif merge_method == 'sum':
+                output = keras.layers.add(sub_networks)
+            elif merge_method == 'concatenate':
+                output = keras.layers.concatenate(sub_networks)
+            else:
+                raise NotImplementedError(
+                    "Unsupported merge method", merge_method)
+
+            result._network = Model(
+                inputs=[peptide_input, allele_input],
+                outputs=[output],
+                name="merged_predictor"
+            )
+            result.update_network_description()
+        else:
+            raise NotImplementedError(
+                "Don't know merge_method to merge networks with layer names: ",
+                layer_names)
+        return result
+
+
     def make_network(
             self,
             peptide_encoding,

test/test_network_merging.py

+from nose.tools import eq_, assert_less, assert_greater, assert_almost_equal, assert_equal
+
+import numpy
+import pandas
+from numpy import testing
+
+numpy.random.seed(0)
+
+import logging
+logging.getLogger('tensorflow').disabled = True
+
+from mhcflurry import Class1AffinityPredictor, Class1NeuralNetwork
+from mhcflurry.common import random_peptides
+from mhcflurry.downloads import get_path
+
+ALLELE_SPECIFIC_PREDICTOR = Class1AffinityPredictor.load(
+    get_path("models_class1", "models"))
+
+PAN_ALLELE_PREDICTOR = Class1AffinityPredictor.load(
+    get_path("models_class1_pan", "models.with_mass_spec"))
+
+
+def test_merge():
+    peptides = random_peptides(100, length=9)
+    peptides.extend(random_peptides(100, length=10))
+    peptides = pandas.Series(peptides).sample(frac=1.0)
+
+    alleles = pandas.Series(
+        ["HLA-A*03:01", "HLA-B*57:01", "HLA-C*02:01"]
+    ).sample(n=len(peptides), replace=True)
+
+    predictions1 = PAN_ALLELE_PREDICTOR.predict(
+        peptides=peptides, alleles=alleles)
+
+    merged = Class1NeuralNetwork.merge(
+        PAN_ALLELE_PREDICTOR.class1_pan_allele_models)
+    merged_predictor = Class1AffinityPredictor(
+        allele_to_sequence=PAN_ALLELE_PREDICTOR.allele_to_sequence,
+        class1_pan_allele_models=[merged],
+    )
+    predictions2 = merged_predictor.predict(peptides=peptides, alleles=alleles)
+    numpy.testing.assert_allclose(predictions1, predictions2, atol=0.1)

test/test_speed.py

@@ -145,6 +145,7 @@ if __name__ == '__main__':
 
     if "pan-allele" in args.predictor:
         print("Running pan-allele test")
+        PAN_ALLELE_PREDICTOR.optimize()
         result = test_speed_pan_allele(
             profile=True, num=args.num_predictions)
         result[