diff --git a/compiler_opt/es/training_state.py b/compiler_opt/es/training_state.py new file mode 100644 index 00000000..22893144 --- /dev/null +++ b/compiler_opt/es/training_state.py @@ -0,0 +1,226 @@ +# Copyright 2020 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Training state management library. + +Manages serialization and parameter history for multi-model evolution strategy +training. Checkpoints are saved in `checkpoint.json` alongside subdirectories +containing the active JAX models and params (in `models.pkl`) and optionally +`.tflite` model files. + +Checkpoint JSON Schema: + step: Current training iteration index. + latest_models_dir: Subdirectory name containing the active model set. + model_keys: List of model identifier strings matching the loaded state. + max_reward: Maximum reward observed so far. + optimizer_state: Internal state configuration of the blackbox optimizer. +""" + +from collections.abc import Callable +import dataclasses +import json +import logging +import os + +from absl import flags +import cloudpickle +from compiler_opt.es import blackbox_optimizers +from compiler_opt.jax import dense_nn_bc +from etils import epath +import jax +import jax.flatten_util +import numpy as np +import numpy.typing as npt + +_INIT_FROM_CHECKPOINT = flags.DEFINE_string( + "init_from_checkpoint", + None, + "Path to a checkpoint to initialize the training state from if no " + "checkpoint exists in the output directory.", +) + +_PICKLED_MODULES_FILENAME = "models.pkl" + + +class TrainingState: + """Container managing historical training parameters and serialization.""" + + def __init__( + self, + output_path: str, + get_initial_models: Callable[ + [], + dict[str, dense_nn_bc.ModelSpec], + ], + optimizer: blackbox_optimizers.BlackboxOptimizer, + ): + self._output_path = epath.Path(output_path) + self._optimizer = optimizer + self._checkpoint_path = self._output_path / "checkpoint.json" + + self._model_specs = {} + self._model_keys = [] + + ckpt = None + ckpt_source_path = None + if self._checkpoint_path.exists(): + logging.info("Loading checkpoint from %s", self._checkpoint_path) + ckpt = json.loads(self._checkpoint_path.read_text()) + ckpt_source_path = self._checkpoint_path + elif _INIT_FROM_CHECKPOINT.value: + logging.info("Loading initial checkpoint from %s", + _INIT_FROM_CHECKPOINT.value) + init_ckpt_path = epath.Path(_INIT_FROM_CHECKPOINT.value) + if not init_ckpt_path.exists(): + raise FileNotFoundError( + f"init_from_checkpoint {_INIT_FROM_CHECKPOINT.value} does not exist" + ) + ckpt = json.loads(init_ckpt_path.read_text()) + ckpt_source_path = init_ckpt_path + + if ckpt is not None: + logging.info("Restoring models from checkpoint %s", ckpt_source_path) + assert ckpt_source_path is not None + ckpt_dir = ckpt_source_path.parent + latest_models_dir = ckpt["latest_models_dir"] + self._model_keys = ckpt["model_keys"] + pickle_path = ckpt_dir / latest_models_dir / _PICKLED_MODULES_FILENAME + with pickle_path.open("rb") as f: + self._model_specs = cloudpickle.load(f) + + self._step = ckpt["step"] + 1 + self._max_reward = ckpt["max_reward"] + self._optimizer.set_state(ckpt["optimizer_state"]) + else: + initial_models = get_initial_models() + if not isinstance(initial_models, dict): + raise ValueError("get_initial_model must return a dict") + self._model_specs = initial_models + self._model_keys = sorted(self._model_specs.keys()) + self._step = 0 + self._max_reward = 0.0 + + self._unravel_fns = {} + self._flat_shapes = {} + for k in self._model_keys: + flat, unravel = jax.flatten_util.ravel_pytree(self._model_specs[k].params) + self._unravel_fns[k] = unravel + self._flat_shapes[k] = flat.shape[0] + + self._model_weights = np.concatenate([ + jax.flatten_util.ravel_pytree(self._model_specs[k].params)[0] + for k in self._model_keys + ]) + + if not self._checkpoint_path.exists(): + init_dir_name = "initial_model" + init_dir = self._output_path / init_dir_name + init_dir.mkdir(parents=True, exist_ok=True) + for k in self._model_keys: + path = os.fspath(init_dir / f"{k}.tflite") + self._model_specs[k].save_tflite(path) + + pickle_path = init_dir / _PICKLED_MODULES_FILENAME + with pickle_path.open("wb") as f: + cloudpickle.dump(self._model_specs, f) + + self._save_checkpoint(self._step - 1, init_dir_name) + + @property + def current_step(self) -> int: + return self._step + + @property + def current_weights(self) -> npt.NDArray[np.float32]: + return self._model_weights + + @property + def model_specs(self) -> dict[str, dense_nn_bc.ModelSpec]: + return self._model_specs + + def save_policies( + self, + rewards: list[float], + perturbations: list[npt.NDArray[np.float32]], + new_model_weights: npt.NDArray[np.float32], + ) -> float: + max_reward = np.max(rewards) + + if self._max_reward < max_reward: + self._max_reward = max_reward + logging.info( + "Found new best model with reward %f at step %d, saving.", + max_reward, + self._step, + ) + max_index = np.argmax(rewards) + perturbation = perturbations[max_index] + + best_weights = self._model_weights + perturbation + best_dir_name = f"best_policy_{max_reward}_{self._step}" + best_dir = self._output_path / best_dir_name + best_dir.mkdir(parents=True, exist_ok=True) + offset = 0 + for k in self._model_keys: + size = self._flat_shapes[k] + model_weights = best_weights[offset:offset + size] + offset += size + + model_params = self._unravel_fns[k](model_weights) + best_model = dense_nn_bc.ModelSpec( + model=self._model_specs[k].model, + params=model_params, + sample_inputs=self._model_specs[k].sample_inputs, + ) + best_model.save_tflite(os.fspath(best_dir / f"{k}.tflite")) + + iter_dir_name = f"iteration_{self._step}" + iter_dir = self._output_path / iter_dir_name + iter_dir.mkdir(parents=True, exist_ok=True) + + offset = 0 + for k in self._model_keys: + size = self._flat_shapes[k] + model_flat_weights = new_model_weights[offset:offset + size] + offset += size + + new_params = self._unravel_fns[k](model_flat_weights) + self._model_specs[k] = dataclasses.replace( + self._model_specs[k], params=new_params) + + self._model_specs[k].save_tflite(os.fspath(iter_dir / f"{k}.tflite")) + + pickle_path = iter_dir / _PICKLED_MODULES_FILENAME + with pickle_path.open("wb") as f: + cloudpickle.dump(self._model_specs, f) + + self._save_checkpoint(self._step, iter_dir_name) + + self._model_weights = new_model_weights + ret_max = self._max_reward + self._step += 1 + + return ret_max + + def _save_checkpoint(self, step: int, latest_models_dir: str) -> None: + ckpt_data = { + "step": step, + "latest_models_dir": latest_models_dir, + "model_keys": self._model_keys, + "max_reward": self._max_reward, + "optimizer_state": self._optimizer.get_state(), + } + + tmp_checkpoint = self._output_path / "checkpoint.json.tmp" + tmp_checkpoint.write_text(json.dumps(ckpt_data, indent=2)) + tmp_checkpoint.replace(self._checkpoint_path) diff --git a/compiler_opt/es/training_state_test.py b/compiler_opt/es/training_state_test.py new file mode 100644 index 00000000..979b9c2b --- /dev/null +++ b/compiler_opt/es/training_state_test.py @@ -0,0 +1,394 @@ +# Copyright 2020 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for training_state management.""" + +import json +import multiprocessing +import pickle +import traceback +from unittest import mock + +from absl.testing import absltest +from absl.testing import flagsaver +from compiler_opt.es import blackbox_optimizers +from compiler_opt.es import training_state +from compiler_opt.jax import dense_nn_bc +from compiler_opt.jax import litert_adapter +from etils import epath +import jax +import jax.flatten_util +import numpy as np + + +def _worker_process_fn(pipe, model_spec, flat_weights, inputs): + try: + pipe.send(("progress", "started")) + + params = jax.tree_util.tree_map(jax.device_put, model_spec.params) + sample_inputs = jax.tree_util.tree_map(jax.device_put, + model_spec.sample_inputs) + flat_weights = jax.device_put(flat_weights) + inputs = jax.tree_util.tree_map(jax.device_put, inputs) + pipe.send(("progress", "converted_to_jax")) + + _, unravel_fn = jax.flatten_util.ravel_pytree(params) + params = unravel_fn(flat_weights) + pipe.send(("progress", "unravelled")) + + edge_model = litert_adapter.convert( + model_spec.model.apply, params, jax_fn_kwargs=sample_inputs) + pipe.send(("progress", "converted")) + + outputs = edge_model(**inputs) + pipe.send(("progress", "evaluated")) + + pipe.send(("success", outputs)) + except Exception as e: # pylint: disable=broad-exception-caught + pipe.send(("failure", f"{str(e)}\n{traceback.format_exc()}")) + finally: + pipe.close() + + +class FakeOptimizer(blackbox_optimizers.BlackboxOptimizer): + + def __init__(self): + self.state = "initial_state" + + def get_state(self): + return self.state + + def set_state(self, state): + self.state = state + + def get_hyperparameters(self): + raise NotImplementedError() + + def run_step(self, current_weights, perturbations, rewards): + raise NotImplementedError() + + def update_state(self, rewards, perturbations): + raise NotImplementedError() + + +def create_dummy_model_spec(num_features=2, hidden_size=4, seed=0): + constants = { + "internal_x_mean": np.zeros(num_features, dtype=np.float32), + "internal_x_std": np.ones(num_features, dtype=np.float32), + } + model = dense_nn_bc.DenseModel( + hidden_layers=[hidden_size], constants=constants) + sample_inputs = {"x": np.zeros((1, num_features), dtype=np.float32)} + key = jax.random.key(seed) + params = model.init(key, **sample_inputs) + return dense_nn_bc.ModelSpec.fromJax( + model=model, params=params, sample_inputs=sample_inputs) + + +class TrainingStateTest(absltest.TestCase): + + def setUp(self): + super().setUp() + self.test_dir = epath.Path(self.create_tempdir().full_path) + self.spec_a = create_dummy_model_spec(2, 4, seed=0) + self.spec_b = create_dummy_model_spec(3, 2, seed=1) + + self.get_initial_models = lambda: { + "model_a": self.spec_a, + "model_b": self.spec_b, + } + self.optimizer = FakeOptimizer() + + @mock.patch.object(litert_adapter, "convert") + def test_constructor_initializes_new_state(self, mock_convert): + mock_edge_model = mock.MagicMock() + mock_edge_model.tflite_model.return_value = b"dummy_tflite_bytes" + mock_convert.return_value = mock_edge_model + + state = training_state.TrainingState( + output_path=str(self.test_dir), + get_initial_models=self.get_initial_models, + optimizer=self.optimizer, + ) + + self.assertEqual(state.current_step, 0) + flat_a, _ = jax.flatten_util.ravel_pytree(self.spec_a.params) + flat_b, _ = jax.flatten_util.ravel_pytree(self.spec_b.params) + expected_weights = np.concatenate([flat_a, flat_b]) + np.testing.assert_array_equal(state.current_weights, expected_weights) + + self.assertTrue( + (self.test_dir / "initial_model" / "model_a.tflite").exists()) + self.assertTrue( + (self.test_dir / "initial_model" / "model_b.tflite").exists()) + + pkl_path = ( + self.test_dir / "initial_model" / + training_state._PICKLED_MODULES_FILENAME) + self.assertTrue(pkl_path.exists()) + with pkl_path.open("rb") as f: + saved_specs = pickle.load(f) + self.assertEqual(saved_specs.keys(), {"model_a", "model_b"}) + np.testing.assert_array_equal( + jax.flatten_util.ravel_pytree(saved_specs["model_a"].params)[0], flat_a) + + checkpoint_path = self.test_dir / "checkpoint.json" + self.assertTrue(checkpoint_path.exists()) + ckpt_data = json.loads(checkpoint_path.read_text()) + self.assertEqual(ckpt_data["step"], -1) + self.assertEqual(ckpt_data["latest_models_dir"], "initial_model") + self.assertEqual(ckpt_data["model_keys"], ["model_a", "model_b"]) + + @mock.patch.object(litert_adapter, "convert") + def test_save_policies_creates_checkpoints_and_best_policy( + self, mock_convert): + mock_edge_model = mock.MagicMock() + mock_edge_model.tflite_model.return_value = b"dummy_tflite_bytes" + mock_convert.return_value = mock_edge_model + + state = training_state.TrainingState( + output_path=str(self.test_dir), + get_initial_models=self.get_initial_models, + optimizer=self.optimizer, + ) + + flat_a, _ = jax.flatten_util.ravel_pytree(self.spec_a.params) + flat_b, _ = jax.flatten_util.ravel_pytree(self.spec_b.params) + size_a = flat_a.shape[0] + size_b = flat_b.shape[0] + total_size = size_a + size_b + + rng = np.random.default_rng(0) + new_weights = rng.random(total_size, dtype=np.float32) + perturbations = [ + rng.random(total_size, dtype=np.float32), + rng.random(total_size, dtype=np.float32), + ] + + max_reward = state.save_policies( + rewards=[10.0, 20.0], + perturbations=perturbations, + new_model_weights=new_weights, + ) + + self.assertEqual(max_reward, 20.0) + self.assertEqual(state.current_step, 1) + np.testing.assert_array_equal(state.current_weights, new_weights) + + best_dir = self.test_dir / "best_policy_20.0_0" + self.assertTrue(best_dir.exists()) + self.assertTrue((best_dir / "model_a.tflite").exists()) + self.assertTrue((best_dir / "model_b.tflite").exists()) + + iter_dir = self.test_dir / "iteration_0" + self.assertTrue(iter_dir.exists()) + self.assertTrue((iter_dir / "model_a.tflite").exists()) + self.assertTrue((iter_dir / "model_b.tflite").exists()) + self.assertTrue( + (iter_dir / training_state._PICKLED_MODULES_FILENAME).exists()) + + checkpoint_path = self.test_dir / "checkpoint.json" + ckpt_data = json.loads(checkpoint_path.read_text()) + self.assertEqual(ckpt_data["step"], 0) + self.assertEqual(ckpt_data["latest_models_dir"], "iteration_0") + self.assertEqual(ckpt_data["max_reward"], 20.0) + + def test_restores_from_checkpoint(self): + checkpoint_data = { + "step": 42, + "latest_models_dir": "iteration_42", + "model_keys": ["model_a", "model_b"], + "max_reward": 99.0, + "optimizer_state": "checkpointed_state", + } + checkpoint_path = self.test_dir / "checkpoint.json" + checkpoint_path.write_text(json.dumps(checkpoint_data)) + + iter_dir = self.test_dir / "iteration_42" + iter_dir.mkdir(parents=True, exist_ok=True) + + restored_spec_a = create_dummy_model_spec(2, 4, seed=2) + restored_spec_b = create_dummy_model_spec(3, 2, seed=3) + + model_specs = { + "model_a": restored_spec_a, + "model_b": restored_spec_b, + } + with (iter_dir / training_state._PICKLED_MODULES_FILENAME).open("wb") as f: + pickle.dump(model_specs, f) + + state = training_state.TrainingState( + output_path=str(self.test_dir), + get_initial_models=self.get_initial_models, + optimizer=self.optimizer, + ) + + self.assertEqual(state.current_step, 43) + self.assertEqual(self.optimizer.get_state(), "checkpointed_state") + + flat_a, _ = jax.flatten_util.ravel_pytree(restored_spec_a.params) + flat_b, _ = jax.flatten_util.ravel_pytree(restored_spec_b.params) + expected_weights = np.concatenate([flat_a, flat_b]) + np.testing.assert_array_equal(state.current_weights, expected_weights) + + @mock.patch.object(litert_adapter, "convert") + def test_restores_from_given_checkpoint(self, mock_convert): + mock_edge_model = mock.MagicMock() + mock_edge_model.tflite_model.return_value = b"dummy_tflite_bytes" + mock_convert.return_value = mock_edge_model + + other_dir = epath.Path(self.create_tempdir().full_path) + checkpoint_data = { + "step": 42, + "latest_models_dir": "iteration_42", + "model_keys": ["model_a", "model_b"], + "max_reward": 99.0, + "optimizer_state": "checkpointed_state", + } + given_checkpoint_path = other_dir / "checkpoint.json" + given_checkpoint_path.write_text(json.dumps(checkpoint_data)) + + iter_dir = other_dir / "iteration_42" + iter_dir.mkdir(parents=True, exist_ok=True) + + restored_spec_a = create_dummy_model_spec(2, 4, seed=2) + restored_spec_b = create_dummy_model_spec(3, 2, seed=3) + model_specs = { + "model_a": restored_spec_a, + "model_b": restored_spec_b, + } + with (iter_dir / training_state._PICKLED_MODULES_FILENAME).open("wb") as f: + pickle.dump(model_specs, f) + + with flagsaver.flagsaver(init_from_checkpoint=str(given_checkpoint_path)): + state = training_state.TrainingState( + output_path=str(self.test_dir), + get_initial_models=self.get_initial_models, + optimizer=self.optimizer, + ) + + self.assertEqual(state.current_step, 43) + self.assertEqual(self.optimizer.get_state(), "checkpointed_state") + + flat_a, _ = jax.flatten_util.ravel_pytree(restored_spec_a.params) + flat_b, _ = jax.flatten_util.ravel_pytree(restored_spec_b.params) + expected_weights = np.concatenate([flat_a, flat_b]) + np.testing.assert_array_equal(state.current_weights, expected_weights) + + self.assertTrue( + (self.test_dir / "initial_model" / "model_a.tflite").exists()) + self.assertTrue( + (self.test_dir / "initial_model" / "model_b.tflite").exists()) + self.assertTrue((self.test_dir / "initial_model" / + training_state._PICKLED_MODULES_FILENAME).exists()) + + def test_worker_simulation_different_process(self): + constants = { + "internal_input_a_mean": np.array([1.0], dtype=np.float32), + "internal_input_a_std": np.array([2.0], dtype=np.float32), + "internal_input_b_mean": np.array([3.0], dtype=np.float32), + } + model = dense_nn_bc.DenseModel(hidden_layers=[2], constants=constants) + sample_inputs = { + "input_a": np.zeros((1, 1), dtype=np.float32), + "input_b": np.zeros((1, 1), dtype=np.float32), + } + key = jax.random.key(0) + params_jax = model.init(key, **sample_inputs) + + params_jax["params"]["Dense_0"]["kernel"] = np.array( + [[1.0, 2.0], [3.0, 4.0]], dtype=np.float32) + params_jax["params"]["Dense_0"]["bias"] = np.array([0.5, -0.5], + dtype=np.float32) + params_jax["params"]["Dense_1"]["kernel"] = np.array([[0.5], [1.5]], + dtype=np.float32) + params_jax["params"]["Dense_1"]["bias"] = np.array([1.0], dtype=np.float32) + + model_spec = dense_nn_bc.ModelSpec( + model=model, params=params_jax, sample_inputs=sample_inputs) + + flat_weights, _ = jax.flatten_util.ravel_pytree(params_jax) + + perturbation = np.arange(1, len(flat_weights) + 1, dtype=np.float32) + + perturbed_weights_numpy = np.array(flat_weights) + perturbation + + def make_inputs(a, b): + return { + "input_a": np.array([[a]], dtype=np.float32), + "input_b": np.array([[b]], dtype=np.float32), + } + + def make_jax_inputs(a, b): + return jax.tree_util.tree_map(jax.device_put, make_inputs(a, b)) + + eval_inputs = make_inputs(5.0, 7.0) + eval_inputs_jax = make_jax_inputs(5.0, 7.0) + + _, unravel_fn = jax.flatten_util.ravel_pytree(params_jax) + perturbed_params = unravel_fn(perturbed_weights_numpy) + perturbed_params_jax = jax.tree_util.tree_map(jax.device_put, + perturbed_params) + expected_output = model.apply(perturbed_params_jax, **eval_inputs_jax) + + variations = [ + expected_output, + model.apply(params_jax, **eval_inputs_jax), + model.apply(params_jax, **make_jax_inputs(-1.0, 3.0)), + model.apply(perturbed_params_jax, **make_jax_inputs(-1.0, 3.0)), + ] + self.assertTrue(all(v.size == 1 for v in variations)) + variations = [v.item() for v in variations] + self.assertEmpty([v for v in variations if np.isnan(v)]) + self.assertLen(set(variations), 4) + + ctx = multiprocessing.get_context("spawn") + parent_conn, child_conn = ctx.Pipe() + p = ctx.Process( + target=_worker_process_fn, + args=(child_conn, model_spec, perturbed_weights_numpy, eval_inputs), + ) + p.start() + + success = False + value = None + while True: + if parent_conn.poll(timeout=120): + msg_type, msg_val = parent_conn.recv() + if msg_type == "progress": + print(f"Worker progress: {msg_val}", flush=True) + elif msg_type == "success": + success = True + value = msg_val + break + elif msg_type == "failure": + self.fail(f"Worker failed with: {msg_val}") + break + else: + p.terminate() + self.fail("Worker process hung (no progress for 120s)") + break + + p.join() + + self.assertTrue(success) + + if isinstance(value, dict): + value = list(value.values())[0] + + expected_output_np = np.array(expected_output) + np.testing.assert_allclose(value, expected_output_np, rtol=1e-5) + + +if __name__ == "__main__": + absltest.main() diff --git a/compiler_opt/jax/__init__.py b/compiler_opt/jax/__init__.py new file mode 100644 index 00000000..6913f02e --- /dev/null +++ b/compiler_opt/jax/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2020 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/compiler_opt/jax/dense_nn_bc.py b/compiler_opt/jax/dense_nn_bc.py new file mode 100644 index 00000000..0c1ec4d5 --- /dev/null +++ b/compiler_opt/jax/dense_nn_bc.py @@ -0,0 +1,350 @@ +# Copyright 2020 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Behavioral cloning for dense neural networks.""" + +from collections.abc import Callable, Iterable, Sequence +import dataclasses +import logging +from typing import Any + +from compiler_opt.jax import litert_adapter +from etils import epath +from flax import linen as nn +import jax +import jax.numpy as jnp +import numpy as np +import optax + +FeatureTransformFn = Callable[[Any], Any] + + +def _identity_fn(x: Any) -> Any: + return x + + +def _default_feature_transforms(*_) -> FeatureTransformFn: + return _identity_fn + + +def _extract_features_and_constants( + records: Iterable[Any], + label_name: str, + feature_transforms: Callable[ + [str, Any], FeatureTransformFn] = _default_feature_transforms, +) -> tuple[dict[str, np.ndarray], np.ndarray, dict[str, np.ndarray]]: + """Extracts raw features, target label, and feature normalization constants.""" + features_accum: dict[str, list[Any]] = {} + + for record in records: + for name, val in record.features.items(): + if name not in features_accum: + features_accum[name] = [] + features_accum[name].append(val.value) + + if not features_accum: + raise ValueError("No data found in records") + + all_features: dict[str, np.ndarray] = {} + for name, values in features_accum.items(): + all_features[name] = np.array(values, dtype=np.float32) + + labels = all_features.pop(label_name) + + constants: dict[str, np.ndarray] = {} + for name, data in all_features.items(): + transformed_data = feature_transforms(name, np)(data) + + mean = np.mean(transformed_data, axis=0) + std = np.std(transformed_data, axis=0) + std = np.where(std == 0, 1.0, std) + constants[f"internal_{name}_mean"] = mean + constants[f"internal_{name}_std"] = std + + return all_features, labels, constants + + +def preprocess_data( + records: Iterable[Any], + label_name: str, + feature_transforms: Callable[ + [str, Any], FeatureTransformFn] = _default_feature_transforms, +) -> tuple[dict[str, np.ndarray], np.ndarray, dict[str, np.ndarray]]: + """Preprocesses records into features, normalized targets, and normalization constants.""" + all_features, target_scores, constants = _extract_features_and_constants( + records, label_name, feature_transforms) + + target_mean = np.mean(target_scores, axis=0, keepdims=True) + target_std = np.std(target_scores, axis=0, keepdims=True) + target_std = np.where(target_std == 0, 1.0, target_std) + + normalized_targets = (target_scores - target_mean) / target_std + return all_features, normalized_targets, constants + + +class DenseModel(nn.Module): + """MLP scoring with baked-in normalization.""" + + hidden_layers: Sequence[int] + constants: dict[str, np.ndarray] + feature_transforms: Callable[[str, Any], FeatureTransformFn] = ( + _default_feature_transforms) + + @nn.compact + def __call__(self, **features): + normalized_features = {} + for k, v in features.items(): + transformed_feature = self.feature_transforms(k, jnp)(v) + mean_key = f"internal_{k}_mean" + std_key = f"internal_{k}_std" + if mean_key in self.constants and std_key in self.constants: + mean = self.constants[mean_key] + std = self.constants[std_key] + normalized_features[k] = (transformed_feature - mean) / std + else: + normalized_features[k] = transformed_feature + + x = jnp.concatenate( + [normalized_features[k] for k in sorted(normalized_features.keys())], + axis=-1, + ) + for feat_size in self.hidden_layers: + x = nn.Dense(features=feat_size)(x) + x = jax.nn.relu(x) + x = nn.Dense(features=1)(x) + return x + + +class DenseClassifierModel(DenseModel): + """MLP binary classifier with baked-in normalization and step decision.""" + + threshold: float = 0.0 + export_binary_decision: bool = False + feature_transforms: Callable[[str, Any], FeatureTransformFn] = ( + _default_feature_transforms) + + @nn.compact + def __call__(self, **features): + logits = super().__call__(**features) + if self.export_binary_decision: + return jnp.where(logits > self.threshold, 1.0, 0.0) + return logits + + +@dataclasses.dataclass(frozen=True) +class ModelSpec: + model: nn.Module + params: dict[str, Any] + sample_inputs: dict[str, Any] + + @staticmethod + def fromJax( + model: nn.Module, + params: dict[str, Any], + sample_inputs: dict[str, Any], + ) -> "ModelSpec": + return ModelSpec( + model=model, + params=jax.tree_util.tree_map(np.array, params), + sample_inputs=jax.tree_util.tree_map(np.array, sample_inputs), + ) + + def save_tflite(self, filename: str): + edge_model = litert_adapter.convert( + self.model.apply, + self.params, + jax_fn_kwargs=self.sample_inputs, + ) + epath.Path(filename).write_bytes(edge_model.tflite_model()) + + +def _train_model( + model: nn.Module, + all_features: dict[str, np.ndarray], + targets: np.ndarray, + loss_fn: Callable[[Any, dict[str, Any], Any], Any], + steps: int, + sample_inputs: dict[str, Any], + batch_size: int = 128, + learning_rate: float = 0.001, + seed: int = 0, + balanced_sampling: bool = False, + pos_indices: np.ndarray | None = None, + neg_indices: np.ndarray | None = None, + loss_name: str = "Loss", + log_prefix: str = "Training", +) -> dict[str, Any]: + """Trains a Flax model on batch features and targets.""" + key = jax.random.key(seed) + params = model.init(key, **sample_inputs) + + grad_fn = jax.value_and_grad(loss_fn) + tx = optax.adam(learning_rate=learning_rate) + opt_state = tx.init(params) + + @jax.jit + def train_step(p, opt_s, feat, target): + loss_val, grads = grad_fn(p, feat, target) + updates, opt_s = tx.update(grads, opt_s, p) + p = optax.apply_updates(p, updates) + return p, opt_s, loss_val + + rng = np.random.default_rng(seed) + num_examples = targets.shape[0] + use_balanced = ( + balanced_sampling and pos_indices is not None and + neg_indices is not None and len(pos_indices) > 0 and len(neg_indices) > 0) + + logging.info("%s for %d steps", log_prefix, steps) + for i in range(steps): + if use_balanced: + pos_batch_size = batch_size // 2 + neg_batch_size = batch_size - pos_batch_size + batch_pos = rng.choice(pos_indices, size=pos_batch_size, replace=True) + batch_neg = rng.choice(neg_indices, size=neg_batch_size, replace=True) + indices = np.concatenate([batch_pos, batch_neg]) + rng.shuffle(indices) + else: + batch_idx = (i * batch_size) % num_examples + indices = np.arange(batch_idx, batch_idx + batch_size) % num_examples + + batch_feat = {k: v[indices] for k, v in all_features.items()} + batch_target = targets[indices] + + params, opt_state, loss_val = train_step(params, opt_state, batch_feat, + batch_target) + if i % 100 == 0: + logging.info("Step %d, %s: %f", i, loss_name, loss_val) + + return params + + +def behavioral_clone( + records: Iterable[Any], + steps: int, + sample_inputs: dict[str, Any], + hidden_layers: Sequence[int], + label_name: str, + batch_size: int = 128, + learning_rate: float = 0.001, + feature_transforms: Callable[ + [str, Any], FeatureTransformFn] = _default_feature_transforms, +) -> ModelSpec: + """Behavioral clones the model from feature records (regression with MSE).""" + all_features, normalized_targets, constants = preprocess_data( + records, label_name, feature_transforms) + + model = DenseModel( + hidden_layers=hidden_layers, + constants=constants, + feature_transforms=feature_transforms, + ) + + def mse_loss(p, feat, target): + predictions = model.apply(p, **feat) + return jnp.mean( + (jnp.squeeze(predictions, axis=-1) - jnp.squeeze(target, axis=-1))**2) + + params = _train_model( + model=model, + all_features=all_features, + targets=normalized_targets, + loss_fn=mse_loss, + steps=steps, + sample_inputs=sample_inputs, + batch_size=batch_size, + learning_rate=learning_rate, + log_prefix="Training", + ) + + return ModelSpec.fromJax( + model=model, params=params, sample_inputs=sample_inputs) + + +def behavioral_clone_classifier( + records: Iterable[Any], + steps: int, + sample_inputs: dict[str, Any], + hidden_layers: Sequence[int], + label_name: str, + batch_size: int = 128, + learning_rate: float = 0.001, + pos_weight: float | None = None, + threshold: float = 0.0, + balanced_sampling: bool = False, + seed: int = 0, + feature_transforms: Callable[ + [str, Any], FeatureTransformFn] = _default_feature_transforms, +) -> ModelSpec: + """Behavioral clones a binary decision classifier from feature records.""" + all_features, labels, constants = _extract_features_and_constants( + records, label_name, feature_transforms) + + train_model = DenseClassifierModel( + hidden_layers=hidden_layers, + constants=constants, + threshold=threshold, + export_binary_decision=False, + feature_transforms=feature_transforms, + ) + + labels = np.array(labels, dtype=np.float32) + if labels.ndim == 1: + labels = np.expand_dims(labels, axis=-1) + + flat_labels = labels.flatten() + pos_indices = np.where(flat_labels > 0.5)[0] + neg_indices = np.where(flat_labels <= 0.5)[0] + logging.info( + "Classifier BC: %d positive labels, %d negative labels", + len(pos_indices), + len(neg_indices), + ) + if pos_weight is None and len(pos_indices) > 0 and len(neg_indices) > 0: + pos_weight = float(len(neg_indices)) / float(len(pos_indices)) + logging.info("Auto-calculated pos_weight: %f", pos_weight) + + def bce_loss(p, feat, target): + logits = train_model.apply(p, **feat) + loss = optax.sigmoid_binary_cross_entropy(logits=logits, labels=target) + if pos_weight is not None: + w = jnp.where(target > 0.5, pos_weight, 1.0) + loss = loss * w + return jnp.mean(loss) + + params = _train_model( + model=train_model, + all_features=all_features, + targets=labels, + loss_fn=bce_loss, + steps=steps, + sample_inputs=sample_inputs, + batch_size=batch_size, + learning_rate=learning_rate, + seed=seed, + balanced_sampling=balanced_sampling, + pos_indices=pos_indices, + neg_indices=neg_indices, + loss_name="BCE Loss", + log_prefix="Training classifier", + ) + + inference_model = DenseClassifierModel( + hidden_layers=hidden_layers, + constants=constants, + threshold=threshold, + export_binary_decision=True, + feature_transforms=feature_transforms, + ) + return ModelSpec.fromJax( + model=inference_model, params=params, sample_inputs=sample_inputs) diff --git a/compiler_opt/jax/dense_nn_bc_test.py b/compiler_opt/jax/dense_nn_bc_test.py new file mode 100644 index 00000000..09393086 --- /dev/null +++ b/compiler_opt/jax/dense_nn_bc_test.py @@ -0,0 +1,362 @@ +# Copyright 2020 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for dense_nn_bc.""" + +from absl.testing import absltest +from compiler_opt.jax import dense_nn_bc +from compiler_opt.jax import litert_adapter +import jax +import jax.numpy as jnp +import numpy as np + + +class DummyRecord: + + def __init__(self, features): + self.features = features + + +class DummyVal: + + def __init__(self, value): + self.value = value + + +class DenseNnBcTest(absltest.TestCase): + + def test_behavioral_clone(self): + rng = np.random.default_rng(0) + records = [] + for _ in range(10): + features = { + "feat1": DummyVal(rng.random(4)), + "feat2": DummyVal(rng.random(2)), + "score": DummyVal(rng.random(1)), + } + records.append(DummyRecord(features)) + + sample_inputs = { + "feat1": jnp.zeros(4), + "feat2": jnp.zeros(2), + } + + spec = dense_nn_bc.behavioral_clone( + records=records, + steps=2, + sample_inputs=sample_inputs, + hidden_layers=[8, 4], + label_name="score", + ) + + self.assertIsNotNone(spec.model) + self.assertIsNotNone(spec.params) + + edge_model = litert_adapter.convert( + spec.model.apply, spec.params, jax_fn_kwargs=spec.sample_inputs) + self.assertIsNotNone(edge_model) + + def test_from_jax(self): + constants = { + "internal_feat1_mean": np.array([0.0]), + "internal_feat1_std": np.array([1.0]), + } + model = dense_nn_bc.DenseModel(hidden_layers=[4], constants=constants) + sample_inputs = {"feat1": jnp.zeros((1, 1))} + key = jax.random.key(0) + params = model.init(key, **sample_inputs) + + spec = dense_nn_bc.ModelSpec.fromJax( + model=model, params=params, sample_inputs=sample_inputs) + + self.assertIsInstance(spec.params["params"]["Dense_0"]["kernel"], + np.ndarray) + self.assertIsInstance(spec.sample_inputs["feat1"], np.ndarray) + + def test_save(self): + constants = { + "internal_feat1_mean": np.array([0.0]), + "internal_feat1_std": np.array([1.0]), + } + model = dense_nn_bc.DenseModel(hidden_layers=[4], constants=constants) + sample_inputs = {"feat1": jnp.zeros((1, 1))} + key = jax.random.key(0) + params = model.init(key, **sample_inputs) + + spec = dense_nn_bc.ModelSpec.fromJax( + model=model, params=params, sample_inputs=sample_inputs) + temp_file = self.create_tempfile() + spec.save_tflite(temp_file.full_path) + self.assertGreater(len(temp_file.read_bytes()), 0) + + def test_preprocess_data_normal(self): + records = [ + DummyRecord({ + "feat1": DummyVal(np.array([1.0, 2.0])), + "score": DummyVal(np.array([10.0])), + }), + DummyRecord({ + "feat1": DummyVal(np.array([3.0, 4.0])), + "score": DummyVal(np.array([20.0])), + }), + ] + features, normalized_targets, constants = dense_nn_bc.preprocess_data( + records, "score") + + self.assertIn("feat1", features) + np.testing.assert_array_almost_equal(features["feat1"], + np.array([[1.0, 2.0], [3.0, 4.0]])) + self.assertNotIn("score", features) + + np.testing.assert_array_almost_equal(normalized_targets, + np.array([[-1.0], [1.0]])) + + np.testing.assert_array_almost_equal(constants["internal_feat1_mean"], + np.array([2.0, 3.0])) + np.testing.assert_array_almost_equal(constants["internal_feat1_std"], + np.array([1.0, 1.0])) + + def test_preprocess_data_zero_variance_target(self): + records = [ + DummyRecord({ + "feat1": DummyVal(np.array([1.0, 2.0])), + "score": DummyVal(np.array([10.0])), + }), + DummyRecord({ + "feat1": DummyVal(np.array([3.0, 4.0])), + "score": DummyVal(np.array([10.0])), + }), + ] + _, normalized_targets, _ = dense_nn_bc.preprocess_data(records, "score") + + np.testing.assert_array_almost_equal(normalized_targets, + np.array([[0.0], [0.0]])) + + def test_preprocess_data_zero_variance_feature(self): + records = [ + DummyRecord({ + "feat1": DummyVal(np.array([1.0, 2.0])), + "score": DummyVal(np.array([10.0])), + }), + DummyRecord({ + "feat1": DummyVal(np.array([1.0, 4.0])), + "score": DummyVal(np.array([20.0])), + }), + ] + _, _, constants = dense_nn_bc.preprocess_data(records, "score") + + np.testing.assert_array_almost_equal(constants["internal_feat1_mean"], + np.array([1.0, 3.0])) + np.testing.assert_array_almost_equal(constants["internal_feat1_std"], + np.array([1.0, 1.0])) + + def test_behavioral_clone_classifier(self): + rng = np.random.default_rng(0) + records = [] + for i in range(10): + label = 1.0 if i < 2 else 0.0 + features = { + "feat1": DummyVal(rng.random(4)), + "feat2": DummyVal(rng.random(2)), + "decision": DummyVal(np.array([label])), + } + records.append(DummyRecord(features)) + + sample_inputs = { + "feat1": jnp.zeros(4), + "feat2": jnp.zeros(2), + } + + spec = dense_nn_bc.behavioral_clone_classifier( + records=records, + steps=5, + sample_inputs=sample_inputs, + hidden_layers=[8, 4], + label_name="decision", + pos_weight=4.0, + balanced_sampling=True, + ) + + self.assertIsNotNone(spec.model) + self.assertIsNotNone(spec.params) + + output = spec.model.apply(spec.params, **sample_inputs) + self.assertIn(float(output[0]), (0.0, 1.0)) + + temp_file = self.create_tempfile() + spec.save_tflite(temp_file.full_path) + self.assertGreater(len(temp_file.read_bytes()), 0) + + def test_behavioral_clone_classifier_auto_pos_weight(self): + rng = np.random.default_rng(0) + records = [] + for i in range(10): + label = 1.0 if i < 2 else 0.0 + features = { + "feat1": DummyVal(rng.random(4)), + "feat2": DummyVal(rng.random(2)), + "decision": DummyVal(np.array([label])), + } + records.append(DummyRecord(features)) + + sample_inputs = { + "feat1": jnp.zeros(4), + "feat2": jnp.zeros(2), + } + + spec = dense_nn_bc.behavioral_clone_classifier( + records=records, + steps=5, + sample_inputs=sample_inputs, + hidden_layers=[8, 4], + label_name="decision", + pos_weight=None, + ) + + self.assertIsNotNone(spec.model) + self.assertIsNotNone(spec.params) + output = spec.model.apply(spec.params, **sample_inputs) + self.assertIn(float(output[0]), (0.0, 1.0)) + + def test_extract_features_and_constants(self): + records = [ + DummyRecord({ + "feat1": DummyVal(np.array([1.0, 2.0])), + "decision": DummyVal(np.array([1.0])), + }), + DummyRecord({ + "feat1": DummyVal(np.array([3.0, 4.0])), + "decision": DummyVal(np.array([0.0])), + }), + ] + features, labels, constants = dense_nn_bc._extract_features_and_constants( + records, "decision") + + self.assertIn("feat1", features) + self.assertNotIn("decision", features) + np.testing.assert_array_equal(labels, np.array([[1.0], [0.0]])) + np.testing.assert_array_almost_equal(constants["internal_feat1_mean"], + np.array([2.0, 3.0])) + np.testing.assert_array_almost_equal(constants["internal_feat1_std"], + np.array([1.0, 1.0])) + + def test_preprocess_data_custom_feature_transforms(self): + + def custom_feature_transforms(name, mod): + if name == "feat1": + return lambda x: mod.exp(x) + return lambda x: x + + records = [ + DummyRecord({ + "feat1": DummyVal(np.array([0.0, 0.0])), + "score": DummyVal(np.array([10.0])), + }), + DummyRecord({ + "feat1": DummyVal(np.array([np.log(3.0), np.log(5.0)])), + "score": DummyVal(np.array([20.0])), + }), + ] + _, _, constants = dense_nn_bc.preprocess_data( + records, "score", feature_transforms=custom_feature_transforms) + + np.testing.assert_array_almost_equal(constants["internal_feat1_mean"], + np.array([2.0, 3.0])) + np.testing.assert_array_almost_equal(constants["internal_feat1_std"], + np.array([1.0, 2.0])) + + def test_behavioral_clone_custom_feature_transforms(self): + + def custom_feature_transforms(name, mod): + if name == "feat1": + return lambda x: mod.exp(x) + return lambda x: x + + records = [ + DummyRecord({ + "feat1": DummyVal(np.array([0.0, 0.0])), + "score": DummyVal(np.array([10.0])), + }), + DummyRecord({ + "feat1": DummyVal(np.array([np.log(3.0), np.log(5.0)])), + "score": DummyVal(np.array([20.0])), + }), + ] + sample_inputs = { + "feat1": jnp.zeros((1, 2)), + } + + spec = dense_nn_bc.behavioral_clone( + records=records, + steps=2, + sample_inputs=sample_inputs, + hidden_layers=[4], + label_name="score", + feature_transforms=custom_feature_transforms, + ) + + self.assertIsNotNone(spec.model) + self.assertIsNotNone(spec.params) + self.assertEqual(spec.model.feature_transforms, custom_feature_transforms) + + output = spec.model.apply(spec.params, **sample_inputs) + self.assertEqual(output.shape, (1, 1)) + + temp_file = self.create_tempfile() + spec.save_tflite(temp_file.full_path) + self.assertGreater(len(temp_file.read_bytes()), 0) + + def test_behavioral_clone_classifier_custom_feature_transforms(self): + + def custom_feature_transforms(name, mod): + if name == "feat1": + return lambda x: mod.exp(x) + return lambda x: x + + records = [ + DummyRecord({ + "feat1": DummyVal(np.array([0.0, 0.0])), + "decision": DummyVal(np.array([1.0])), + }), + DummyRecord({ + "feat1": DummyVal(np.array([np.log(3.0), np.log(5.0)])), + "decision": DummyVal(np.array([0.0])), + }), + ] + sample_inputs = { + "feat1": jnp.zeros((1, 2)), + } + + spec = dense_nn_bc.behavioral_clone_classifier( + records=records, + steps=2, + sample_inputs=sample_inputs, + hidden_layers=[4], + label_name="decision", + feature_transforms=custom_feature_transforms, + ) + + self.assertIsNotNone(spec.model) + self.assertIsNotNone(spec.params) + self.assertEqual(spec.model.feature_transforms, custom_feature_transforms) + + output = spec.model.apply(spec.params, **sample_inputs) + self.assertIn(float(output[0, 0]), (0.0, 1.0)) + + temp_file = self.create_tempfile() + spec.save_tflite(temp_file.full_path) + self.assertGreater(len(temp_file.read_bytes()), 0) + + +if __name__ == "__main__": + absltest.main() diff --git a/compiler_opt/jax/litert_adapter.py b/compiler_opt/jax/litert_adapter.py new file mode 100644 index 00000000..f0b2b50c --- /dev/null +++ b/compiler_opt/jax/litert_adapter.py @@ -0,0 +1,140 @@ +# Copyright 2020 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Open-source LiteRT conversion and execution adapter for JAX models.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from typing import Any + +from ai_edge_litert.interpreter import Interpreter +from etils import epath +from jax.experimental import jax2tf +import numpy as np +import tensorflow as tf + + +class TfLiteModel: + """Lightweight wrapper around serialized LiteRT flatbuffer bytes.""" + + def __init__(self, model_bytes: bytes): + self._model_bytes = bytes(model_bytes) + self._interpreter: Interpreter | None = None + + def tflite_model(self) -> bytes: + return self._model_bytes + + def _get_interpreter(self) -> Interpreter: + if self._interpreter is None: + self._interpreter = Interpreter(model_content=self._model_bytes) + self._interpreter.allocate_tensors() + return self._interpreter + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + interp = self._get_interpreter() + if kwargs: + runner = interp.get_signature_runner("serving_default") + np_kwargs = {k: np.asarray(v, dtype=np.float32) for k, v in kwargs.items()} + outputs = runner(**np_kwargs) + return outputs + + input_details = interp.get_input_details() + if len(args) != len(input_details): + raise ValueError( + f"Expected {len(input_details)} positional inputs, got {len(args)}") + for detail, arg in zip(input_details, args): + interp.set_tensor(detail["index"], np.asarray(arg, dtype=detail["dtype"])) + interp.invoke() + output_details = interp.get_output_details() + outputs = [interp.get_tensor(detail["index"]) for detail in output_details] + return outputs[0] if len(outputs) == 1 else outputs + + +class Model: + """Loader matching the internal LiteRT model interface.""" + + @staticmethod + def load(path: str | epath.Path) -> TfLiteModel: + return TfLiteModel(epath.Path(path).read_bytes()) + + +def get_signature_list(edge_model: TfLiteModel) -> dict[str, Any]: + """Returns the signature list dictionary for a loaded LiteRT model.""" + return edge_model._get_interpreter().get_signature_list() # pylint: disable=protected-access + + +def convert( + jax_fn: Callable[..., Any], + params: Any, + jax_fn_args: Sequence[Any] = (), + jax_fn_kwargs: Mapping[str, Any] | None = None, +) -> TfLiteModel: + """Converts a JAX callable and params into a TfLiteModel.""" + kwargs_dict = dict(jax_fn_kwargs) if jax_fn_kwargs is not None else {} + + if kwargs_dict and jax_fn_args: + raise ValueError("Specify either jax_fn_args or jax_fn_kwargs, not both.") + + if kwargs_dict: + kw_keys = sorted(kwargs_dict.keys()) + specs = [ + tf.TensorSpec( + shape=np.asarray(kwargs_dict[k]).shape, + dtype=tf.dtypes.as_dtype(np.asarray(kwargs_dict[k]).dtype), + name=k, + ) for k in kw_keys + ] + + def _fn_pos(*pos_args): + bound = dict(zip(kw_keys, pos_args)) + return jax_fn(params, **bound) + + tf_fn = jax2tf.convert(_fn_pos, with_gradient=False) + + class _Module(tf.Module): + + @tf.function(input_signature=specs) + def serving_default(self, *pos_args): + return {"output_0": tf_fn(*pos_args)} + + mod = _Module() + concrete_fn = mod.serving_default.get_concrete_function(*specs) + else: + args_list = [np.asarray(a) for a in jax_fn_args] + specs = [ + tf.TensorSpec( + shape=a.shape, + dtype=tf.dtypes.as_dtype(a.dtype), + name=f"arg_{i}", + ) for i, a in enumerate(args_list) + ] + + def _fn_args(*pos_args): + return jax_fn(params, *pos_args) + + tf_fn = jax2tf.convert(_fn_args, with_gradient=False) + + class _Module(tf.Module): + + @tf.function(input_signature=specs) + def serving_default(self, *pos_args): + return tf_fn(*pos_args) + + mod = _Module() + concrete_fn = mod.serving_default.get_concrete_function(*specs) + + converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_fn], + mod) + tflite_bytes = converter.convert() + return TfLiteModel(tflite_bytes) diff --git a/compiler_opt/jax/litert_es.py b/compiler_opt/jax/litert_es.py new file mode 100644 index 00000000..831a41da --- /dev/null +++ b/compiler_opt/jax/litert_es.py @@ -0,0 +1,215 @@ +# Copyright 2020 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Library for LiteRT evolutionary strategy APIs.""" + +from collections.abc import Generator +from typing import Any + +from ai_edge_litert import schema_py_generated as model_schema +from ai_edge_litert.tools import flatbuffer_utils +from compiler_opt.jax import litert_adapter +from etils import epath +import numpy as np +import numpy.typing as npt + +_CURRIED_CONSTANT_BUFFERS_META_KEY = "curried_constant_buffers" + + +def _tensor_type_to_np_dtype(tensor_type: int) -> npt.DTypeLike: + """Maps TFLite tensor type enum to numpy dtype.""" + if tensor_type == model_schema.TensorType.FLOAT32: + return np.float32 + if tensor_type == model_schema.TensorType.INT32: + return np.int32 + if tensor_type == model_schema.TensorType.INT8: + return np.int8 + if tensor_type == model_schema.TensorType.UINT8: + return np.uint8 + if tensor_type == model_schema.TensorType.FLOAT16: + return np.float16 + if tensor_type == model_schema.TensorType.INT64: + return np.int64 + if tensor_type == model_schema.TensorType.BOOL: + return np.bool_ + raise ValueError(f"Unsupported tensor type {tensor_type}") + + +def _get_indices_of_ignored_buffers(model_t: model_schema.ModelT) -> set[int]: + """Returns the set of buffer indices that represent curried constants.""" + if not model_t.metadata: + return set() + for meta in model_t.metadata: + meta_name = ( + meta.name.decode("utf-8") + if isinstance(meta.name, bytes) else meta.name) + if meta_name == _CURRIED_CONSTANT_BUFFERS_META_KEY: + buf_idx = meta.buffer + if buf_idx < 0 or buf_idx >= len(model_t.buffers): + raise ValueError(f"Invalid metadata buffer index {buf_idx}") + buf = model_t.buffers[buf_idx] + if buf.data is None or len(buf.data) == 0: + raise ValueError( + f"Metadata buffer {buf_idx} for {_CURRIED_CONSTANT_BUFFERS_META_KEY}" + " is empty") + return set(np.frombuffer(buf.data, dtype=np.int32).tolist()) + return set() + + +def _get_weight_tensors( + model_t: model_schema.ModelT, +) -> Generator[ + tuple[model_schema.TensorT, model_schema.BufferT, type[Any], np.ndarray], + None, + None, +]: + """Yields weight tensors, buffers, dtypes, and numpy arrays.""" + ignored_buffers = _get_indices_of_ignored_buffers(model_t) + for subgraph in model_t.subgraphs: + for tensor in subgraph.tensors: + if tensor.buffer == 0: + # the input + continue + if tensor.buffer in ignored_buffers: + continue + if tensor.buffer >= len(model_t.buffers): + raise ValueError(f"Invalid buffer index {tensor.buffer}") + buffer = model_t.buffers[tensor.buffer] + if buffer.data is not None and len(buffer.data) > 0: + dtype = _tensor_type_to_np_dtype(tensor.type) + arr = np.frombuffer(buffer.data, dtype=dtype) + yield tensor, buffer, dtype, arr + + +def get_weights(model: litert_adapter.TfLiteModel) -> npt.NDArray[np.float32]: + """Extracts weights from a loaded LiteRT model into a numpy array.""" + model_bytes = model.tflite_model() + + model_t = flatbuffer_utils.convert_bytearray_to_object(model_bytes) + + weight_buffers = [] + for _, _, _, arr in _get_weight_tensors(model_t): + weight_buffers.append(arr.flatten()) + + return (np.concatenate(weight_buffers) + if weight_buffers else np.array([], dtype=np.float32)) + + +def save_with_weights( + model: litert_adapter.TfLiteModel, + weights: npt.NDArray[np.float32], + path: str, +) -> None: + """Saves a LiteRT model with updated weights to a file.""" + model_bytes = model.tflite_model() + + model_t = flatbuffer_utils.convert_bytearray_to_object(model_bytes) + + weight_tensors = [] + total_size = 0 + for tensor, buffer, dtype, arr in _get_weight_tensors(model_t): + weight_tensors.append((tensor, buffer, dtype, arr.size)) + total_size += arr.size + + if weights.size != total_size: + raise ValueError( + f"Weight size mismatch. Expected {total_size}, got {weights.size}") + + offset = 0 + for _, buffer, dtype, size in weight_tensors: + chunk = weights[offset:offset + size] + offset += size + new_data = chunk.astype(dtype).tobytes() + buffer.data = bytearray(new_data) + + new_bytes = flatbuffer_utils.convert_object_to_bytearray(model_t) + + epath.Path(path).write_bytes(bytes(new_bytes)) + + +def curry_constant_inputs( + model: litert_adapter.TfLiteModel, + constants: dict[str, np.ndarray], +) -> litert_adapter.TfLiteModel: + """Replaces named inputs with constant buffers and tags them with metadata.""" + model_bytes = model.tflite_model() + model_t = flatbuffer_utils.convert_bytearray_to_object(model_bytes) + + updated_tensors = set() + new_constant_buffers = [] + matched_constants = set() + + for signature in model_t.signatureDefs: + subgraph_idx = signature.subgraphIndex + subgraph = model_t.subgraphs[subgraph_idx] + + new_sig_inputs = [] + inputs_to_remove_from_subgraph = set() + + for tensor_map in signature.inputs: + name = ( + tensor_map.name.decode("utf-8") + if isinstance(tensor_map.name, bytes) else tensor_map.name) + if name in constants: + matched_constants.add(name) + tensor_idx = tensor_map.tensorIndex + tensor = subgraph.tensors[tensor_idx] + + if tensor_idx not in updated_tensors: + val = constants[name] + dtype = _tensor_type_to_np_dtype(tensor.type) + val = np.asarray(val, dtype=dtype) + + new_buffer = model_schema.BufferT() + new_buffer.data = bytearray(val.tobytes()) + model_t.buffers.append(new_buffer) + new_buf_idx = len(model_t.buffers) - 1 + tensor.buffer = new_buf_idx + + new_constant_buffers.append(new_buf_idx) + updated_tensors.add(tensor_idx) + + inputs_to_remove_from_subgraph.add(tensor_idx) + else: + new_sig_inputs.append(tensor_map) + + signature.inputs = new_sig_inputs + + new_subgraph_inputs = [ + idx for idx in subgraph.inputs + if idx not in inputs_to_remove_from_subgraph + ] + subgraph.inputs = new_subgraph_inputs + + missing = set(constants.keys()) - matched_constants + if missing: + raise ValueError( + f"Constant inputs not found in model signatures: {sorted(missing)}") + + if new_constant_buffers: + meta_buf_data = np.array(new_constant_buffers, dtype=np.int32).tobytes() + meta_buffer = model_schema.BufferT() + meta_buffer.data = bytearray(meta_buf_data) + model_t.buffers.append(meta_buffer) + meta_buf_idx = len(model_t.buffers) - 1 + + if model_t.metadata is None: + model_t.metadata = [] + + meta_entry = model_schema.MetadataT() + meta_entry.name = _CURRIED_CONSTANT_BUFFERS_META_KEY + meta_entry.buffer = meta_buf_idx + model_t.metadata.append(meta_entry) + + new_bytes = bytes(flatbuffer_utils.convert_object_to_bytearray(model_t)) + return litert_adapter.TfLiteModel(new_bytes) diff --git a/compiler_opt/jax/litert_es_test.py b/compiler_opt/jax/litert_es_test.py new file mode 100644 index 00000000..b85c1d8d --- /dev/null +++ b/compiler_opt/jax/litert_es_test.py @@ -0,0 +1,112 @@ +# Copyright 2020 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for litert_es.""" + +import os + +from absl.testing import absltest +from compiler_opt.jax import litert_adapter +from compiler_opt.jax import litert_es +import jax.numpy as jnp +import numpy as np + + +class LiteRtEsTest(absltest.TestCase): + + def test_flow(self): + + def model_fn(params, x): + return jnp.dot(x, params["w"]) + params["b"] + + params = { + "w": jnp.array([[2.0], [3.0]], dtype=jnp.float32), + "b": jnp.array([1.0], dtype=jnp.float32), + } + x = jnp.array([[1.0, 2.0]], dtype=jnp.float32) + + edge_model = litert_adapter.convert( + jax_fn=model_fn, params=params, jax_fn_args=(x,)) + + weights = litert_es.get_weights(edge_model) + self.assertEqual(weights.size, 3) + + orig_output = edge_model(x) + if isinstance(orig_output, list): + orig_output = orig_output[0] + np.testing.assert_allclose(orig_output, np.array([[9.0]]), rtol=1e-5) + + doubled_weights = weights * 2 + + new_model_path = os.path.join(self.create_tempdir().full_path, + "new_model.tflite") + litert_es.save_with_weights(edge_model, doubled_weights, new_model_path) + + new_model = litert_adapter.Model.load(new_model_path) + + new_output = new_model(x) + if isinstance(new_output, list): + new_output = new_output[0] + np.testing.assert_allclose(new_output, np.array([[18.0]]), rtol=1e-5) + + def test_curry_constant_inputs(self): + + def model_fn(params, **kwargs): + return jnp.dot(kwargs["x"], params["w"]) + params["b"] + kwargs["c"] + + params = { + "w": jnp.array([[2.0], [3.0]], dtype=jnp.float32), + "b": jnp.array([1.0], dtype=jnp.float32), + } + x = jnp.array([[1.0, 2.0]], dtype=jnp.float32) + c = jnp.array([[5.0]], dtype=jnp.float32) + + edge_model = litert_adapter.convert( + jax_fn=model_fn, params=params, jax_fn_kwargs={ + "x": x, + "c": c + }) + + curried_model = litert_es.curry_constant_inputs( + edge_model, {"c": np.array([[5.0]], dtype=np.float32)}) + + sig_list = litert_adapter.get_signature_list(curried_model) + self.assertIn("serving_default", sig_list) + self.assertEqual(list(sig_list["serving_default"]["inputs"]), ["x"]) + + out = curried_model(x=x) + if isinstance(out, dict): + out = list(out.values())[0] + elif isinstance(out, list): + out = out[0] + np.testing.assert_allclose(out, np.array([[14.0]]), rtol=1e-5) + + weights = litert_es.get_weights(curried_model) + self.assertEqual(weights.size, 3) + + doubled_weights = weights * 2 + new_model_path = os.path.join(self.create_tempdir().full_path, + "curried_new_model.tflite") + litert_es.save_with_weights(curried_model, doubled_weights, new_model_path) + + new_model = litert_adapter.Model.load(new_model_path) + new_out = new_model(x=x) + if isinstance(new_out, dict): + new_out = list(new_out.values())[0] + elif isinstance(new_out, list): + new_out = new_out[0] + np.testing.assert_allclose(new_out, np.array([[23.0]]), rtol=1e-5) + + +if __name__ == "__main__": + absltest.main() diff --git a/run_jax_tests.sh b/run_jax_tests.sh index 32a52267..68e5ed81 100755 --- a/run_jax_tests.sh +++ b/run_jax_tests.sh @@ -24,6 +24,9 @@ TEST_TARGETS=( "compiler_opt/distributed/worker_test.py" "compiler_opt/distributed/buffered_scheduler_test.py" "compiler_opt/distributed/local/local_worker_manager_test.py" + "compiler_opt/jax/dense_nn_bc_test.py" + "compiler_opt/jax/litert_es_test.py" + "compiler_opt/es/training_state_test.py" ) python3 -m pytest "${TEST_TARGETS[@]}"