diff --git a/deepvariant/make_examples_options.py b/deepvariant/make_examples_options.py index f84a2c75..72870142 100644 --- a/deepvariant/make_examples_options.py +++ b/deepvariant/make_examples_options.py @@ -790,7 +790,9 @@ _SMALL_MODEL_INFERENCE_BATCH_SIZE = flags.DEFINE_integer( 'small_model_inference_batch_size', 128, - 'Sets the batch size used by the small model during inference.', + 'Deprecated and ignored: the small model now runs each region in a single' + ' forward pass, so this batch size no longer affects inference. Retained' + ' for command-line compatibility.', ) _SMALL_MODEL_VAF_CONTEXT_WINDOW_SIZE = flags.DEFINE_integer( 'small_model_vaf_context_window_size', diff --git a/deepvariant/small_model/inference.py b/deepvariant/small_model/inference.py index f8ac47b1..917ebe00 100644 --- a/deepvariant/small_model/inference.py +++ b/deepvariant/small_model/inference.py @@ -69,7 +69,17 @@ def passes_confidence_threshold( def classify( classifier: keras.Model, examples: np.ndarray, batch_size: int ) -> np.ndarray: - return classifier.predict(examples, batch_size=batch_size, verbose=0) + # batch_size (the deprecated --small_model_inference_batch_size) is ignored: + # each region's candidates are run in a single forward pass. + del batch_size + # Call the model directly rather than via Model.predict(). predict() rebuilds + # a data adapter, callback list, and prediction loop on every invocation; + # because the small model is called once per region on a small batch, that + # per-call scaffolding -- not the arithmetic -- dominates its runtime. + # __call__ runs the same (cached) forward pass without that overhead. The + # result is numerically identical here: the model is a plain Dense MLP with no + # train/inference-divergent layers (e.g. BatchNorm or Dropout). + return np.asarray(classifier(examples, training=False)) class SmallModelVariantCaller: diff --git a/deepvariant/small_model/inference_test.py b/deepvariant/small_model/inference_test.py index eead9592..c4cf8da6 100644 --- a/deepvariant/small_model/inference_test.py +++ b/deepvariant/small_model/inference_test.py @@ -29,6 +29,7 @@ from unittest import mock from absl.testing import absltest from absl.testing import parameterized +import numpy as np import tensorflow as tf from deepvariant.protos import deepvariant_pb2 from deepvariant.small_model import inference @@ -85,7 +86,9 @@ class SmallModelVariantCallerTest(parameterized.TestCase): def setUp(self): super().setUp() self.mock_classifier = mock.MagicMock() - self.mock_classifier.predict.return_value = [ + # `classify` invokes the model directly (classifier(...)), so stub __call__ + # rather than .predict(). + self.mock_classifier.return_value = [ (0.0, 0.999, 0.0), (0.0, 0.0, 0.1), (0.0, 0.0, 0.999), @@ -192,7 +195,7 @@ def test_call_variants_small_batch(self): with mock.patch.object( type(self.mock_classifier), "__call__", - lambda self, x, training: [ + lambda self, x, training=False: [ (0.0, 0.999, 0.0), (0.0, 0.0, 0.1), ], @@ -256,5 +259,57 @@ def test_emit_all_candidates_enabled(self): self.assertLen(candidates_not_called, 4) +class ClassifyEquivalenceTest(parameterized.TestCase): + """`classify` calls the model directly; that must match Model.predict().""" + + def _build_small_model(self, num_features: int = 70): + """Builds a minimal Dense + softmax MLP of the small model's general shape. + + The predict()-vs-__call__ equivalence holds for any layer sizes (it is a + property of the forward pass, not the architecture), so a small network is + used for speed rather than the production layer widths. Weights are then + amplified so the softmax is peaked, making the arg-max decision meaningful + instead of noise from near-uniform untrained outputs. + """ + keras = inference.keras + model = keras.Sequential([ + keras.layers.Input(shape=(num_features,)), + keras.layers.Dense(32, activation="relu"), + keras.layers.Dense(16, activation="relu"), + keras.layers.Dense( + len(make_small_model_examples.GenotypeEncoding), + activation="softmax", + ), + ]) + for layer in model.layers: + weights = layer.get_weights() + if weights: + layer.set_weights([w * 5.0 for w in weights]) + return model + + @parameterized.parameters(1, 2, 7, 64, 128, 300) + def test_classify_matches_predict(self, num_candidates: int): + num_features = 70 + model = self._build_small_model(num_features) + rng = np.random.default_rng(num_candidates) + # Integer features, exactly as encode_inference_examples produces them. + examples = rng.integers( + 0, 100, size=(num_candidates, num_features) + ).astype(np.int64) + + expected = model.predict(examples, batch_size=128, verbose=0) + actual = inference.classify(model, examples, batch_size=128) + + self.assertEqual(actual.shape, expected.shape) + self.assertEqual(actual.dtype, expected.dtype) + # classify() runs the identical forward graph as predict() (same weights, no + # train/inference-divergent layers), so the shape, dtype, and every genotype + # arg-max decision are preserved; the probabilities also match to sub-ULP. + np.testing.assert_allclose(actual, expected, rtol=1e-6, atol=1e-6) + np.testing.assert_array_equal( + np.argmax(actual, axis=1), np.argmax(expected, axis=1) + ) + + if __name__ == "__main__": absltest.main()