diff --git a/.gitattributes b/.gitattributes index 61287f73..05b5f9f8 100644 --- a/.gitattributes +++ b/.gitattributes @@ -1,3 +1,4 @@ * text=auto eol=lf *.zip binary *.npy binary +*.py eol=lf diff=python diff --git a/Makefile b/Makefile index aa19e9e6..327dc41d 100644 --- a/Makefile +++ b/Makefile @@ -7,6 +7,9 @@ sample_model: unet2d: cd tests/data/unet2d && zip -r $(ROOT_DIR)/unet2d.tmodel ./* +dummy_tf: + cd tests/data/dummy_tensorflow && zip -r $(ROOT_DIR)/dummy_tf.tmodel ./* + protos: python -m grpc_tools.protoc -I./proto --python_out=tiktorch/proto/ --grpc_python_out=tiktorch/proto/ ./proto/*.proto sed -i -r 's/import (.+_pb2.*)/from . import \1/g' tiktorch/proto/*_pb2*.py @@ -27,4 +30,4 @@ remove_devenv: conda env remove --yes --name $(TIKTORCH_ENV_NAME) -.PHONY: protos version sample_model devenv remove_devenv +.PHONY: protos version sample_model devenv remove_devenv dummy_tf diff --git a/tests/conftest.py b/tests/conftest.py index f7e1c447..7844d491 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -17,6 +17,7 @@ TEST_DATA = "data" TEST_PYBIO_ZIPFOLDER = "unet2d" TEST_PYBIO_DUMMY = "dummy" +TEST_PYBIO_TENSORFLOW_DUMMY = "dummy_tensorflow" NNModel = namedtuple("NNModel", ["model", "state"]) @@ -114,6 +115,34 @@ def pybio_dummy_model_bytes(data_path): return data +def archive(directory): + result = io.BytesIO() + + with ZipFile(result, mode="w") as zip_model: + + def _archive(path_to_archive): + for path in path_to_archive.iterdir(): + if str(path.name).startswith("__"): + continue + + if path.is_dir(): + _archive(path) + + else: + with path.open(mode="rb") as f: + zip_model.writestr(str(path).replace(str(directory), ""), f.read()) + + _archive(directory) + + return result + + +@pytest.fixture +def pybio_dummy_tensorflow_model_bytes(data_path): + pybio_net_dir = Path(data_path) / TEST_PYBIO_TENSORFLOW_DUMMY + return archive(pybio_net_dir) + + @pytest.fixture def cache_path(tmp_path): return Path(getenv("PYBIO_CACHE_PATH", tmp_path)) diff --git a/tests/data/dummy_tensorflow/Dummy.model.yaml b/tests/data/dummy_tensorflow/Dummy.model.yaml new file mode 100644 index 00000000..7ea71669 --- /dev/null +++ b/tests/data/dummy_tensorflow/Dummy.model.yaml @@ -0,0 +1,43 @@ +name: DummyTFModel +description: A dummy tensorflow model for testing +authors: + - ilastik team +cite: + - text: "Ilastik" + doi: https://doi.org +documentation: dummy.md +tags: [tensorflow] +license: MIT + +format_version: 0.1.0 +language: python +framework: tensorflow + +source: dummy.py::TensorflowModelWrapper + +test_input: null # ../test_input.npy +test_output: null # ../test_output.npy + +# TODO double check inputs/outputs +inputs: + - name: input + axes: cyx + data_type: float32 + data_range: [-inf, inf] + shape: [1, 128, 128] +outputs: + - name: output + axes: bcyx + data_type: float32 + data_range: [0, 1] + shape: + reference_input: input # FIXME(m-novikov) ignoring for now + scale: [1, 1, 1] + offset: [0, 0, 0] + #halo: [0, 0, 32, 32] # Should be moved to outputs + +prediction: + weights: + source: ./model + hash: {md5: TODO} + dependencies: conda:./environment.yaml diff --git a/tests/data/dummy_tensorflow/dummy.py b/tests/data/dummy_tensorflow/dummy.py new file mode 100644 index 00000000..d0d6fce3 --- /dev/null +++ b/tests/data/dummy_tensorflow/dummy.py @@ -0,0 +1,12 @@ +class TensorflowModelWrapper: + def __init__(self): + self._model = None + + def set_model(self, model): + self._model = model + + def forward(self, input_): + return self._model.predict(input_) + + def __call__(self, *args, **kwargs): + return self._model.predict(*args, **kwargs) diff --git a/tests/data/dummy_tensorflow/model/saved_model.pb b/tests/data/dummy_tensorflow/model/saved_model.pb new file mode 100644 index 00000000..6b0b590e Binary files /dev/null and b/tests/data/dummy_tensorflow/model/saved_model.pb differ diff --git a/tests/data/dummy_tensorflow/model/variables/variables.data-00000-of-00001 b/tests/data/dummy_tensorflow/model/variables/variables.data-00000-of-00001 new file mode 100644 index 00000000..05278660 Binary files /dev/null and b/tests/data/dummy_tensorflow/model/variables/variables.data-00000-of-00001 differ diff --git a/tests/data/dummy_tensorflow/model/variables/variables.index b/tests/data/dummy_tensorflow/model/variables/variables.index new file mode 100644 index 00000000..202aa363 Binary files /dev/null and b/tests/data/dummy_tensorflow/model/variables/variables.index differ diff --git a/tests/test_server/test_exemplum.py b/tests/test_server/test_exemplum.py index 3afcd0fc..c5319eeb 100644 --- a/tests/test_server/test_exemplum.py +++ b/tests/test_server/test_exemplum.py @@ -1,10 +1,8 @@ -from pathlib import Path - import numpy import torch from pybio.spec import load_spec_and_kwargs -from tiktorch.server.exemplum import Exemplum +from tiktorch.server.model_adapter._exemplum import Exemplum def test_exemplum(data_path, cache_path): @@ -12,9 +10,9 @@ def test_exemplum(data_path, cache_path): assert spec_path.exists(), spec_path.absolute() pybio_model = load_spec_and_kwargs(str(spec_path), cache_path=cache_path) - exemplum = Exemplum(pybio_model=pybio_model, _devices=[torch.device("cpu")]) + exemplum = Exemplum(pybio_model=pybio_model, devices=[torch.device("cpu")]) test_ipt = numpy.load(pybio_model.spec.test_input) # test input with batch dim - out = exemplum.forward(test_ipt[0]) # todo: exemplum.forward should get batch with batch dim + out = exemplum.forward(test_ipt[0].astype(numpy.float32)) # todo: exemplum.forward should get batch with batch dim # assert isinstance(out_seq, (list, tuple)) # todo: forward should return a list # assert len(out_seq) == 1 # out = out_seq diff --git a/tests/test_server/test_reader.py b/tests/test_server/test_reader.py index 614462f2..c020e9f0 100644 --- a/tests/test_server/test_reader.py +++ b/tests/test_server/test_reader.py @@ -3,7 +3,7 @@ import pytest -from tiktorch.server.exemplum import Exemplum +from tiktorch.server.model_adapter import ModelAdapter from tiktorch.server.reader import eval_model_zip, guess_model_path @@ -20,4 +20,10 @@ def test_guess_model_path_without_model_file(paths): def test_eval_model_zip(pybio_model_bytes, cache_path): with ZipFile(pybio_model_bytes) as zf: exemplum = eval_model_zip(zf, devices=["cpu"], cache_path=cache_path) - assert isinstance(exemplum, Exemplum) + assert isinstance(exemplum, ModelAdapter) + + +def test_eval_tensorflow_model_zip(pybio_dummy_tensorflow_model_bytes, cache_path): + with ZipFile(pybio_dummy_tensorflow_model_bytes) as zf: + exemplum = eval_model_zip(zf, devices=["cpu"], cache_path=cache_path) + assert isinstance(exemplum, ModelAdapter) diff --git a/tests/test_server/test_training/test_training.py b/tests/test_server/test_training/test_training.py index 6d5bc3c3..4742022b 100644 --- a/tests/test_server/test_training/test_training.py +++ b/tests/test_server/test_training/test_training.py @@ -27,7 +27,7 @@ def set_break_callback(self, cb): self._break_cb = cb def forward(self, input_tensor): - return torch.Tensor([42]) + return np.array([42]) def set_max_num_iterations(self, val): self.max_num_iterations = val diff --git a/tiktorch/server/model_adapter/__init__.py b/tiktorch/server/model_adapter/__init__.py new file mode 100644 index 00000000..f343e13f --- /dev/null +++ b/tiktorch/server/model_adapter/__init__.py @@ -0,0 +1,21 @@ +from typing import List + +from pybio.spec import nodes + +from ._base import ModelAdapter + +__all__ = ["ModelAdapter", "create_model_adapter"] + + +def create_model_adapter(*, pybio_model: nodes.Model, devices=List[str]): + spec = pybio_model.spec + if spec.framework == "pytorch": + from ._exemplum import Exemplum + + return Exemplum(pybio_model=pybio_model, devices=devices) + elif spec.framework == "tensorflow": + from ._tensorflow_model_adapter import TensorflowModelAdapter + + return TensorflowModelAdapter(pybio_model=pybio_model, devices=devices) + else: + raise NotImplementedError(f"Unknown framework: {spec.framework}") diff --git a/tiktorch/server/model_adapter/_base.py b/tiktorch/server/model_adapter/_base.py new file mode 100644 index 00000000..5f24d5af --- /dev/null +++ b/tiktorch/server/model_adapter/_base.py @@ -0,0 +1,26 @@ +import abc +from typing import Callable + + +class ModelAdapter(abc.ABC): + @abc.abstractmethod + def forward(self, input_tensor): + ... + + @property + @abc.abstractmethod + def max_num_iterations(self) -> int: + ... + + @property + @abc.abstractmethod + def iteration_count(self) -> int: + ... + + @abc.abstractmethod + def set_break_callback(self, thunk: Callable[[], bool]) -> None: + ... + + @abc.abstractmethod + def set_max_num_iterations(self, val: int) -> None: + ... diff --git a/tiktorch/server/exemplum.py b/tiktorch/server/model_adapter/_exemplum.py similarity index 71% rename from tiktorch/server/exemplum.py rename to tiktorch/server/model_adapter/_exemplum.py index fe58b25c..837d7f66 100644 --- a/tiktorch/server/exemplum.py +++ b/tiktorch/server/model_adapter/_exemplum.py @@ -6,14 +6,10 @@ from pybio.spec import nodes from pybio.spec.utils import get_instance -logger = logging.getLogger(__name__) -# @dataclass -# class ValidationOutput(IterationOutput): -# pass +from ._base import ModelAdapter +from ._utils import has_batch_dim -# @dataclass -# class TrainingOutput(IterationOutput): -# pass +logger = logging.getLogger(__name__) def _noop(tensor): @@ -28,29 +24,15 @@ def _add_batch_dim(tensor): return tensor.reshape((1,) + tensor.shape) -def _check_batch_dim(axes: str) -> bool: - try: - index = axes.index("b") - except ValueError: - return False - else: - if index != 0: - raise ValueError("Batch dimension is only supported in first position") - return True - - -class Exemplum: +class Exemplum(ModelAdapter): def __init__( self, *, pybio_model: nodes.Model, - batch_size: int = 1, - num_iterations_per_update: int = 2, - _devices=Sequence[torch.device], + devices=Sequence[str], ): - self.max_num_iterations = 0 - self.iteration_count = 0 - self.devices = _devices + self._max_num_iterations = 0 + self._iteration_count = 0 spec = pybio_model.spec self.name = spec.name @@ -65,7 +47,7 @@ def __init__( self._internal_input_axes = _input.axes self._internal_output_axes = _output.axes - if _check_batch_dim(self._internal_input_axes): + if has_batch_dim(self._internal_input_axes): self.input_axes = self._internal_input_axes[1:] self._input_batch_dimension_transform = _add_batch_dim _input_shape = _input.shape[1:] @@ -78,7 +60,7 @@ def __init__( _halo = _output.halo or [0 for _ in _output.axes] - if _check_batch_dim(self._internal_output_axes): + if has_batch_dim(self._internal_output_axes): self.output_axes = self._internal_output_axes[1:] self._output_batch_dimension_transform = _remove_batch_dim _halo = _halo[1:] @@ -89,29 +71,34 @@ def __init__( self.halo = list(zip(self.output_axes, _halo)) self.model = get_instance(pybio_model) - self.model.to(self.devices[0]) if spec.framework == "pytorch": + self.devices = [torch.device(d) for d in devices] + self.model.to(self.devices[0]) assert isinstance(self.model, torch.nn.Module) if spec.prediction.weights is not None: state = torch.load(spec.prediction.weights.source, map_location=self.devices[0]) self.model.load_state_dict(state) + # elif spec.framework == "tensorflow": + # import tensorflow as tf + # self.devices = [] + # tf_model = tf.keras.models.load_model(spec.prediction.weights.source) + # self.model.set_model(tf_model) else: raise NotImplementedError self._prediction_preprocess = make_concatenated_apply([get_instance(tf) for tf in spec.prediction.preprocess]) self._prediction_postprocess = make_concatenated_apply([get_instance(tf) for tf in spec.prediction.postprocess]) - # inference_engine = ignite.engine.Engine(self._inference_step_function) - # .add_event_handler(Events.STARTED, self.prepare_engine) - # .add_event_handler(Events.COMPLETED, self.log_compute_time) - # def _validation_step_function(self) -> ValidationOutput: - # return ValidationOutput() - # - # - # def _training_step_function(self) -> TrainingOutput: - # return TrainingOutput() + @property + def max_num_iterations(self) -> int: + return self._max_num_iterations + + @property + def iteration_count(self) -> int: + return self._iteration_count def forward(self, batch) -> List[Any]: + batch = torch.from_numpy(batch) with torch.no_grad(): batch = self._input_batch_dimension_transform(batch) batch = self._prediction_preprocess(batch) @@ -120,10 +107,14 @@ def forward(self, batch) -> List[Any]: batch = self._prediction_postprocess(batch) batch = self._output_batch_dimension_transform(batch) assert all([bs > 0 for bs in batch[0].shape]), batch[0].shape - return batch[0] + result = batch[0] + if isinstance(result, torch.Tensor): + return result.detach().cpu().numpy() + else: + return result def set_max_num_iterations(self, max_num_iterations: int) -> None: - self.max_num_iterations = max_num_iterations + self._max_num_iterations = max_num_iterations def set_break_callback(self, cb): return NotImplementedError diff --git a/tiktorch/server/model_adapter/_tensorflow_model_adapter.py b/tiktorch/server/model_adapter/_tensorflow_model_adapter.py new file mode 100644 index 00000000..7089e4c3 --- /dev/null +++ b/tiktorch/server/model_adapter/_tensorflow_model_adapter.py @@ -0,0 +1,98 @@ +from typing import Callable, List + +import numpy as np +import tensorflow as tf +from pybio.spec import nodes +from pybio.spec.utils import get_instance + +from ._base import ModelAdapter +from ._utils import has_batch_dim + + +def _noop(tensor): + return tensor + + +def _remove_batch_dim(batch: List): + return [t.reshape(t.shape[1:]) for t in batch] + + +def _add_batch_dim(tensor): + return tensor.reshape((1,) + tensor.shape) + + +class TensorflowModelAdapter(ModelAdapter): + def __init__( + self, + *, + pybio_model: nodes.Model, + devices=List[str], + ): + spec = pybio_model.spec + self.name = spec.name + + if len(spec.inputs) != 1 or len(spec.outputs) != 1: + raise NotImplementedError("Only single input, single output models are supported") + + assert len(spec.inputs) == 1 + assert len(spec.outputs) == 1 + assert spec.framework == "tensorflow" + + _input = spec.inputs[0] + _output = spec.outputs[0] + + # FIXME: TF probably uses different axis names + self._internal_input_axes = _input.axes + self._internal_output_axes = _output.axes + + if has_batch_dim(self._internal_input_axes): + self.input_axes = self._internal_input_axes[1:] + self._input_batch_dimension_transform = _add_batch_dim + _input_shape = _input.shape[1:] + else: + self.input_axes = self._internal_input_axes + self._input_batch_dimension_transform = _noop + _input_shape = _input.shape + + self.input_shape = list(zip(self.input_axes, _input_shape)) + + _halo = _output.halo or [0 for _ in _output.axes] + + if has_batch_dim(self._internal_output_axes): + self.output_axes = self._internal_output_axes[1:] + self._output_batch_dimension_transform = _remove_batch_dim + _halo = _halo[1:] + else: + self.output_axes = self._internal_output_axes + self._output_batch_dimension_transform = _noop + + self.halo = list(zip(self.output_axes, _halo)) + + self.model = get_instance(pybio_model) + self.devices = [] + tf_model = tf.keras.models.load_model(spec.prediction.weights.source) + self.model.set_model(tf_model) + + def forward(self, input_tensor): + tf_tensor = tf.convert_to_tensor(input_tensor) + + res = self.model.forward(tf_tensor) + + if isinstance(res, np.ndarray): + return res + else: + return tf.make_ndarray(res) + + @property + def max_num_iterations(self) -> int: + return 0 + + @property + def iteration_count(self) -> int: + return 0 + + def set_break_callback(self, thunk: Callable[[], bool]) -> None: + pass + + def set_max_num_iterations(self, val: int) -> None: + pass diff --git a/tiktorch/server/model_adapter/_utils.py b/tiktorch/server/model_adapter/_utils.py new file mode 100644 index 00000000..3bdc6764 --- /dev/null +++ b/tiktorch/server/model_adapter/_utils.py @@ -0,0 +1,9 @@ +def has_batch_dim(axes: str) -> bool: + try: + index = axes.index("b") + except ValueError: + return False + else: + if index != 0: + raise ValueError("Batch dimension is only supported in first position") + return True diff --git a/tiktorch/server/reader.py b/tiktorch/server/reader.py index 1f7b52c1..da9cab6b 100644 --- a/tiktorch/server/reader.py +++ b/tiktorch/server/reader.py @@ -5,11 +5,10 @@ from typing import List, Optional, Sequence from zipfile import ZipFile -import torch from pybio import spec from pybio.spec.utils import train -from tiktorch.server.exemplum import Exemplum +from tiktorch.server.model_adapter import ModelAdapter, create_model_adapter MODEL_EXTENSIONS = (".model.yaml", ".model.yml") logger = logging.getLogger(__name__) @@ -23,8 +22,8 @@ def guess_model_path(file_names: List[str]) -> Optional[str]: return None -def eval_model_zip(model_zip: ZipFile, devices: Sequence[str], cache_path: Optional[Path] = None): - temp_path = Path(tempfile.mkdtemp(prefix="tiktorch")) +def eval_model_zip(model_zip: ZipFile, devices: Sequence[str], cache_path: Optional[Path] = None) -> ModelAdapter: + temp_path = Path(tempfile.mkdtemp(prefix="tiktorch_")) if cache_path is None: cache_path = temp_path / "cache" @@ -38,16 +37,14 @@ def eval_model_zip(model_zip: ZipFile, devices: Sequence[str], cache_path: Optio pybio_model = spec.utils.load_model(spec_file_str, root_path=temp_path, cache_path=cache_path) - devices = [torch.device(d) for d in devices] if pybio_model.spec.training is None: - return Exemplum(pybio_model=pybio_model, _devices=devices) + return create_model_adapter(pybio_model=pybio_model, devices=devices) else: ret = train(pybio_model, _devices=devices) - assert isinstance(ret, Exemplum) - def _on_errror(function, path, exc_info): + def _on_error(function, path, exc_info): logger.warning("Failed to delete temp directory %s", path) - shutil.rmtree(temp_path, on_error=_on_errror) + shutil.rmtree(temp_path, onerror=_on_error) return ret diff --git a/tiktorch/server/session/backend/base.py b/tiktorch/server/session/backend/base.py index f80c9523..45570057 100644 --- a/tiktorch/server/session/backend/base.py +++ b/tiktorch/server/session/backend/base.py @@ -6,7 +6,7 @@ from concurrent.futures import Future from tiktorch.configkeys import TRAINING, VALIDATION -from tiktorch.server.exemplum import Exemplum +from tiktorch.server.model_adapter import ModelAdapter from tiktorch.server.session import types from tiktorch.server.session.backend import commands, supervisor from tiktorch.tiktypes import TikTensorBatch @@ -15,7 +15,7 @@ class SessionBackend: - def __init__(self, exemplum: Exemplum): + def __init__(self, exemplum: ModelAdapter): self._supervisor = supervisor.Supervisor(exemplum) self._supervisor_thread = threading.Thread(target=self._supervisor.run, name="ModelThread") self._supervisor_thread.start() diff --git a/tiktorch/server/session/backend/supervisor.py b/tiktorch/server/session/backend/supervisor.py index 0f11f01c..bf39b242 100644 --- a/tiktorch/server/session/backend/supervisor.py +++ b/tiktorch/server/session/backend/supervisor.py @@ -3,9 +3,9 @@ import logging import queue -import torch +import numpy as np -from tiktorch.server.exemplum import Exemplum +from tiktorch.server.model_adapter import ModelAdapter from tiktorch.server.session import types from tiktorch.server.session.backend import commands @@ -13,7 +13,7 @@ class Supervisor: - def __init__(self, exemplum: Exemplum) -> None: + def __init__(self, exemplum: ModelAdapter) -> None: self._state = types.State.Stopped self._command_queue = commands.CommandPriorityQueue() @@ -39,12 +39,9 @@ def has_work(self): return self._exemplum.max_num_iterations and self._exemplum.max_num_iterations > self._exemplum.iteration_count def forward(self, input_tensor): - torch_input = torch.from_numpy(input_tensor) - result = self._exemplum.forward(torch_input) - if isinstance(result, torch.Tensor): - return result.detach().cpu().numpy() - else: - return result + result = self._exemplum.forward(input_tensor) + assert isinstance(result, np.ndarray) + return result def transition_to(self, new_state: types.State) -> None: logger.debug("Attempting transition to state %s", new_state) diff --git a/vendor/python-bioimage-io b/vendor/python-bioimage-io index 580c70c4..ec230bcc 160000 --- a/vendor/python-bioimage-io +++ b/vendor/python-bioimage-io @@ -1 +1 @@ -Subproject commit 580c70c4b2079f5334a6ad1017287e6a3e4ba895 +Subproject commit ec230bcc6c0d3194b3bb87bb5763f6ff628d3d6b