From b294ccb64013a9ede75483c5b01a79659cadfad1 Mon Sep 17 00:00:00 2001 From: dhritinaidu Date: Thu, 11 Jul 2024 16:28:52 -0400 Subject: [PATCH 01/11] updated requirements --- requirements.txt | 2 ++ 1 file changed, 2 insertions(+) diff --git a/requirements.txt b/requirements.txt index 2f13312..7c8365c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,7 @@ viam-sdk +typing-extensions numpy pyinstaller google-api-python-client +torchvision torch==2.2.1 From 46f8250ddc0abb0de982fe208a00361325622136 Mon Sep 17 00:00:00 2001 From: dhritinaidu Date: Fri, 12 Jul 2024 12:14:19 -0400 Subject: [PATCH 02/11] added workflows and unit tests --- .github/workflows/publish.yml | 18 +++ .github/workflows/test.yml | 25 ++++ .gitignore | 25 ++++ Makefile | 12 ++ build.sh | 23 ++++ meta.json | 18 +++ requirements.txt | 4 +- run.sh | 9 +- src/main.py | 6 +- src/model/model.py | 55 +++++--- src/model_inspector/input_size_calculator.py | 129 +++++++++++------- src/model_inspector/input_tester.py | 23 +++- src/model_inspector/inspector.py | 43 ++++-- .../test_input_size_calculator.py | 97 ------------- src/model_inspector/test_input_tester.py | 60 -------- src/model_inspector/utils.py | 20 ++- src/torch_mlmodel_module.py | 48 +++++-- test_local.py => tests/test_local.py | 43 +++++- 18 files changed, 390 insertions(+), 268 deletions(-) create mode 100644 .github/workflows/publish.yml create mode 100644 .github/workflows/test.yml create mode 100644 .gitignore create mode 100644 Makefile create mode 100755 build.sh create mode 100644 meta.json delete mode 100644 src/model_inspector/test_input_size_calculator.py delete mode 100644 src/model_inspector/test_input_tester.py rename test_local.py => tests/test_local.py (69%) diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml new file mode 100644 index 0000000..1235597 --- /dev/null +++ b/.github/workflows/publish.yml @@ -0,0 +1,18 @@ +on: + push: + tags: + - '[0-9]+.[0-9]+.[0-9]+' + - '[0-9]+.[0-9]+.[0-9]+-rc' + +jobs: + publish: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v3 + - uses: viamrobotics/build-action@v4 + with: + # note: you can replace this line with 'version: ""' if you want to test the build process without deploying + version: ${{ github.ref_name }} + ref: ${{ github.sha }} + key-id: ${{ secrets.viam_key_id }} + key-value: ${{ secrets.viam_key_value }} \ No newline at end of file diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml new file mode 100644 index 0000000..b5d5fe3 --- /dev/null +++ b/.github/workflows/test.yml @@ -0,0 +1,25 @@ +name: Run lint & unit tests + +on: + push: + branches: [ "main" ] + pull_request: + branches: [ "main" ] + +jobs: + build: + name: "Run unit tests" + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + + - name: Install dependencies + run: pip install -r requirements.txt + + - name: Run lint + run: | + make lint + + - name: Run unit tests + run: make test \ No newline at end of file diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..a7ee162 --- /dev/null +++ b/.gitignore @@ -0,0 +1,25 @@ + +__pycache__ +src/model_inspector/__pycache__ +src/model/__pycache__ + +.venv + +build + +dist + +lib + +main.spec +pyvenv.cfg +bin/ + +tests/.pytest_cache + +.pytest_cache + +.DS_Store + + + diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..31f9b93 --- /dev/null +++ b/Makefile @@ -0,0 +1,12 @@ +PYTHONPATH := ./torch:$(PYTHONPATH) + +.PHONY: test lint setup dist + +test: + PYTHONPATH=$(PYTHONPATH) pytest tests/ +lint: + pylint --disable=E1101,W0719,C0202,R0801,W0613 src/ +setup: + python3 -m pip install -r requirements.txt -U +dist/archive.tar.gz: + tar -czf module.tar.gz run.sh requirements.txt src \ No newline at end of file diff --git a/build.sh b/build.sh new file mode 100755 index 0000000..51cb2b2 --- /dev/null +++ b/build.sh @@ -0,0 +1,23 @@ +#!/bin/bash +set -e +UNAME=$(uname -s) +if [ "$UNAME" = "Linux" ] +then + if dpkg -l python3-venv; then + echo "python3-venv is installed, skipping setup" + else + echo "Installing venv on Linux" + sudo apt-get install -y python3-venv + fi +fi +if [ "$UNAME" = "Darwin" ] +then +echo "Installing venv on Darwin" +brew install virtualenv +fi +source .env +python3 -m venv .venv . +source .venv/bin/activate +pip3 install -r requirements.txt +python3 -m PyInstaller --onefile --hidden-import="googleapiclient" src/main.py +tar -czvf dist/archive.tar.gz dist/main \ No newline at end of file diff --git a/meta.json b/meta.json new file mode 100644 index 0000000..1d681c6 --- /dev/null +++ b/meta.json @@ -0,0 +1,18 @@ +{ + "module_id": "viam:torch-cpu", + "visibility": "public", + "url": "https://github.com/viam-labs/torch", + "description": "Viam ML Module service serving PyTorch models.", + "models": [ + { + "api": "rdk:service:mlmodel", + "model": "viam:mlmodel:torch-cpu" + } + ], + "build": { + "build": "./build.sh", + "path": "dist/archive.tar.gz", + "arch": ["linux/arm64", "linux/amd64"] + }, + "entrypoint": "dist/main" +} \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 7c8365c..23e856d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,7 @@ viam-sdk typing-extensions -numpy +numpy==1.24.1 pyinstaller google-api-python-client torchvision -torch==2.2.1 +torch==2.2.1 \ No newline at end of file diff --git a/run.sh b/run.sh index 59956a3..c30ed20 100755 --- a/run.sh +++ b/run.sh @@ -5,9 +5,10 @@ set -euxo pipefail cd $(dirname $0) exec dist/main $@ -# source .env -# ./setup.sh - +source .env +./setup.sh +source .venv/bin/activate # # Be sure to use `exec` so that termination signals reach the python process, # # or handle forwarding termination signals manually -# exec $PYTHON -m src.main $@ +echo which python3 +python3 -m src.main $@ diff --git a/src/main.py b/src/main.py index 8b4f76c..c1cca5e 100644 --- a/src/main.py +++ b/src/main.py @@ -1,11 +1,14 @@ +"""Main runner method of the module""" + import asyncio from viam.module.module import Module from viam.resource.registry import Registry, ResourceCreatorRegistration -from torch_mlmodel_module import TorchMLModelModule from viam.services.mlmodel import MLModel +from torch_mlmodel_module import TorchMLModelModule + async def main(): """ @@ -27,4 +30,5 @@ async def main(): if __name__ == "__main__": + asyncio.run(main()) diff --git a/src/model/model.py b/src/model/model.py index c5ae9af..69fb950 100644 --- a/src/model/model.py +++ b/src/model/model.py @@ -1,53 +1,74 @@ -import torch +""" +This module provides a class for loading and performing inference with a PyTorch model. +The TorchModel class handles loading a serialized model, preparing inputs, and wrapping outputs. +""" + +import os from typing import List, Iterable, Dict, Any -from numpy.typing import NDArray -import torch.nn as nn from collections import OrderedDict + +from numpy.typing import NDArray from viam.logging import getLogger -import os + +import torch +from torch import nn + LOGGER = getLogger(__name__) class TorchModel: + """ + A class to load a PyTorch model from a serialized file or use a provided model, + prepare inputs for the model, perform inference, and wrap the outputs. + """ + def __init__( self, path_to_serialized_file: str, model: nn.Module = None, ) -> None: + "Initializes the model by loading it from a serialized file or using a provided model." if model is not None: self.model = model else: - sizeMB = os.stat(path_to_serialized_file).st_size / (1024 * 1024) - if sizeMB > 500: + size_mb = os.stat(path_to_serialized_file).st_size / (1024 * 1024) + if size_mb > 500: + # pylint: disable=deprecated-method LOGGER.warn( - "model file may be large for certain hardware (" - + str(sizeMB) - + "MB)" + "model file may be large for certain hardware (%s MB)", + size_mb ) self.model = torch.load(path_to_serialized_file) if not isinstance(self.model, nn.Module): if isinstance(self.model, OrderedDict): LOGGER.error( - f"the file {path_to_serialized_file} provided as model file is of type collections.OrderedDict, which suggests that the provided file describes weights instead of a standalone model" + """the file %s provided as model file + is of type collections.OrderedDict, + which suggests that the provided file + describes weights instead of a standalone model""", + path_to_serialized_file ) raise TypeError( f"the model is of type {type(self.model)} instead of nn.Module type" ) self.model.eval() - def infer(self, input): - input = self.prepare_input(input) + def infer(self, input_data): + "Prepares the input data, performs inference using the model, and wraps the output." + input_data = self.prepare_input(input_data) with torch.no_grad(): - output = self.model(*input) + output = self.model(*input_data) return self.wrap_output(output) @staticmethod def prepare_input(input_tensor: Dict[str, NDArray]) -> List[NDArray]: + "Converts a dictionary of NumPy arrays into a list of PyTorch tensors." return [torch.from_numpy(tensor) for tensor in input_tensor.values()] @staticmethod def wrap_output(output: Any) -> Dict[str, NDArray]: + "Converts the output from a PyTorch model to a dictionary of NumPy arrays." if isinstance(output, Iterable): if len(output) == 1: output = output[0] # unpack batched results @@ -55,13 +76,14 @@ def wrap_output(output: Any) -> Dict[str, NDArray]: if isinstance(output, torch.Tensor): return {"output_0": output.numpy()} - elif isinstance(output, dict): + if isinstance(output, dict): for tensor_name, tensor in output.items(): if isinstance(tensor, torch.Tensor): output[tensor_name] = tensor.numpy() return output - elif isinstance(output, Iterable): + + if isinstance(output, Iterable): res = {} count = 0 for out in output: @@ -69,5 +91,4 @@ def wrap_output(output: Any) -> Dict[str, NDArray]: count += 1 return res - else: - raise TypeError(f"can't convert output of type {type(output)} to array") + raise TypeError(f"can't convert output of type {type(output)} to array") diff --git a/src/model_inspector/input_size_calculator.py b/src/model_inspector/input_size_calculator.py index 48799d8..e094591 100644 --- a/src/model_inspector/input_size_calculator.py +++ b/src/model_inspector/input_size_calculator.py @@ -1,7 +1,17 @@ -from model_inspector.utils import is_defined_shape -import torch.nn as nn +""" +InputSizeCalculator module provides methods to calculate +the input size for various types of neural network layers. + +Given a layer and its output shape, this module determines the +input size by leveraging specific methods for different layer types. +It includes methods for linear, RNN, LSTM, +embedding, normalization, pooling, and convolutional layers. +""" from typing import Dict, Tuple from viam.logging import getLogger +from torch import nn +from src.model_inspector.utils import is_defined_shape + LOGGER = getLogger(__name__) @@ -32,6 +42,7 @@ class InputSizeCalculator: def linear( layer: nn.Linear, output_shape: Tuple[int], return_all: bool = False ) -> Tuple[int]: + "Calculates the input size for a linear (fully connected) layer." return ( 1, layer.in_features, @@ -41,32 +52,35 @@ def linear( def rnn( layer: nn.RNN, output_shape: Tuple[int], return_all: bool = False ) -> Tuple[int]: - H_in = layer.input_size + "Calculates the input size for a recurrent neural network (RNN) layer." + h_in = layer.input_size if output_shape is None: - L = -1 + l = -1 elif layer.batch_first: - L = output_shape[1] + l = output_shape[1] else: - L = output_shape[0] - return (L, H_in) + l = output_shape[0] + return (l, h_in) @staticmethod def lstm( layer: nn.LSTM, output_shape: Tuple[int], return_all: bool = False ) -> Tuple[int]: - H_in = layer.input_size + "Calculates the input size for a long short-term memory (LSTM) layer." + h_in = layer.input_size if output_shape is None: - L = -1 + l = -1 elif layer.batch_first: - L = output_shape[1] + l = output_shape[1] else: - L = output_shape[0] - return (L, H_in) + l = output_shape[0] + return (l, h_in) @staticmethod def embedding( layer: nn.Embedding, output_shape: Tuple[int], return_all: bool = False ) -> Tuple[int]: + "Calculates the input size for an embedding layer." if output_shape is None: return -1 return output_shape[0] @@ -75,67 +89,71 @@ def embedding( def layer_norm( layer: nn.LayerNorm, output_shape: Tuple[int], return_all: bool = False ) -> Tuple[int]: + "Returns the normalized shape of a layer normalization (LayerNorm) layer." return layer.normalized_shape @staticmethod def batch_norm_1d( layer: nn.BatchNorm1d, output_shape: Tuple[int], return_all: bool = False ) -> Tuple[int]: - C = layer.num_features + "Calculates the input size for a 1D batch normalization (BatchNorm1d) layer." + c = layer.num_features if output_shape is None: - L = -1 + l = -1 else: - L = output_shape[-1] - return (C, L) + l = output_shape[-1] + return (c, l) @staticmethod def batch_norm_2d( layer: nn.BatchNorm2d, output_shape: Tuple[int], return_all: bool = False ) -> Tuple[int]: - C = layer.num_features + "Calculates the input size for a 2D batch normalization (BatchNorm2d) layer." + c = layer.num_features if output_shape is None: - H, W = -1, -1 + h, w = -1, -1 else: - H, W = output_shape[-2], output_shape[-1] - return (C, H, W) + h, w = output_shape[-2], output_shape[-1] + return (c, h, w) @staticmethod def batch_norm_3d( layer: nn.BatchNorm3d, output_shape: Tuple[int], return_all: bool = False ) -> Tuple[int]: - # TODO + "Calculates the input size for a 3D batch normalization (BatchNorm3d) layer." + return output_shape @staticmethod def maxpool_1d( layer: nn.MaxPool1d, output_shape: Tuple[int], return_all: bool = False ) -> Tuple[int]: + "Calculates the input size for a 1D max pooling (MaxPool1d) layer." if output_shape is None or not is_defined_shape(output_shape): return (-1, -1) - C, L_out = output_shape + c, l_out = output_shape padding = layer.padding dilation = layer.dilation ks = layer.kernel_size stride = layer.stride - L_in = stride * (L_out - 1) - 2 * padding + dilation * (ks - 1) + 1 + l_in = stride * (l_out - 1) - 2 * padding + dilation * (ks - 1) + 1 if return_all: res = [] for i in range(stride): - res.append((C, L_in + i)) + res.append((c, l_in + i)) return res - - else: - return (C, L_in) + return (c, l_in) @staticmethod def maxpool_2d( layer: nn.MaxPool2d, output_shape: Tuple[int], return_all: bool = False ) -> Tuple[int]: + "Calculates the input size for a 2D max pooling (MaxPool2d) layer." if output_shape is None or not is_defined_shape(output_shape): return (-1, -1, -1) # (C, H, W) - C_out, H_out, W_out = output_shape + c_out, h_out, w_out = output_shape padding = layer.padding dilation = layer.dilation ks = layer.kernel_size @@ -146,102 +164,103 @@ def maxpool_2d( ks = layer.kernel_size stride = layer.stride - H_in = stride[0] * (H_out - 1) - 2 * padding[0] + dilation[0] * (ks[0] - 1) + 1 - W_in = stride[1] * (W_out - 1) - 2 * padding[1] + dilation[1] * (ks[1] - 1) + 1 + h_in = stride[0] * (h_out - 1) - 2 * padding[0] + dilation[0] * (ks[0] - 1) + 1 + w_in = stride[1] * (w_out - 1) - 2 * padding[1] + dilation[1] * (ks[1] - 1) + 1 if return_all: res = [] for i in range(stride[0]): for j in range(stride[1]): - res.append((C_out, H_in + i, W_in + j)) + res.append((c_out, h_in + i, w_in + j)) return res - else: - return (C_out, H_in, W_in) + return (c_out, h_in, w_in) @staticmethod def avgpool_1d( layer: nn.AvgPool1d, output_shape: Tuple[int], return_all: bool = False ) -> Tuple[int]: + "Calculates the input size for a 1D average pooling (AvgPool1d) layer." if output_shape is None or not is_defined_shape(output_shape): return (-1, -1) - C, L_out = output_shape + c, l_out = output_shape padding = layer.padding ks = layer.kernel_size stride = layer.stride - L_in = stride * (L_out - 1) - 2 * padding + ks + l_in = stride * (l_out - 1) - 2 * padding + ks if return_all: res = [] for i in range(stride): - res.append((C, L_in + i)) + res.append((c, l_in + i)) return res - else: - return (C, L_in) + return (c, l_in) @staticmethod def avgpool_2d( layer: nn.AvgPool2d, output_shape: Tuple[int], return_all: bool = False ) -> Tuple[int]: + "Calculates the input size for a 2D average pooling (AvgPool2d) layer." if output_shape is None or not is_defined_shape(output_shape): return (-1, -1, -1) - C, L_out = output_shape + c, l_out = output_shape padding = layer.padding ks = layer.kernel_size stride = layer.stride - L_in = stride * (L_out - 1) - 2 * padding + ks + l_in = stride * (l_out - 1) - 2 * padding + ks if return_all: res = [] for i in range(stride): - res.append((C, L_in + i)) + res.append((c, l_in + i)) return res - else: - return (C, L_in) + return (c, l_in) @staticmethod def conv_1d( layer: nn.Conv1d, output_shape: Tuple[int], return_all: bool = False ) -> Tuple[int]: + "Calculates the input size for a 1D convolutional (Conv1d) layer." if output_shape is None or not is_defined_shape(output_shape): return (layer.in_channels, -1) - C, L_out = output_shape + c, l_out = output_shape padding = layer.padding dilation = layer.dilation ks = layer.kernel_size stride = layer.stride - L_in = stride * (L_out - 1) - 2 * padding + dilation * (ks - 1) + 1 + l_in = stride * (l_out - 1) - 2 * padding + dilation * (ks - 1) + 1 res = [] for i in range(stride): - res.append((C, L_in + i)) + res.append((c, l_in + i)) return res @staticmethod def conv_2d( layer: nn.Conv2d, output_shape: Tuple[int], return_all: bool = False ) -> Tuple[int]: + """Calculates the input size for a given 2D convolutional layer + based on the provided output shape.""" if output_shape is None or not is_defined_shape(output_shape): return (layer.in_channels, -1, -1) - H_out, W_out = output_shape[-2], output_shape[-1] + h_out, w_out = output_shape[-2], output_shape[-1] padding = layer.padding dilation = layer.dilation ks = layer.kernel_size stride = layer.stride - H_in = stride[0] * (H_out - 1) - 2 * padding[0] + dilation[0] * (ks[0] - 1) + 1 - W_in = stride[1] * (W_out - 1) - 2 * padding[1] + dilation[1] * (ks[1] - 1) + 1 + h_in = stride[0] * (h_out - 1) - 2 * padding[0] + dilation[0] * (ks[0] - 1) + 1 + w_in = stride[1] * (w_out - 1) - 2 * padding[1] + dilation[1] * (ks[1] - 1) + 1 res = [] if return_all: for i in range(stride[0]): for j in range(stride[1]): - res.append((layer.in_channels, H_in + i, W_in + j)) + res.append((layer.in_channels, h_in + i, w_in + j)) return res - else: - return (layer.in_channels, H_in, W_in) + return (layer.in_channels, h_in, w_in) @staticmethod def default(layer, output_shape, return_all: bool = False): @@ -284,11 +303,15 @@ def __init__(self): def get_input_size( self, layer: nn.Module, output_shape: Tuple[int], return_all: bool = False ) -> Tuple[int]: + "Calculates the input size for a given layer based on the provided output shape." layer_type = type(layer) calculator = self.calculators.get(layer_type, self.default) input_shape = calculator(layer, output_shape, return_all) LOGGER.info( - f" For layer type {type(layer).__name__}, with output shape: {output_shape} input shape found is {input_shape}" + "For layer type %s, with output shape: %s input shape found is %s", + type(layer).__name__, + output_shape, + input_shape ) return input_shape diff --git a/src/model_inspector/input_tester.py b/src/model_inspector/input_tester.py index bfaaf17..051cbab 100644 --- a/src/model_inspector/input_tester.py +++ b/src/model_inspector/input_tester.py @@ -1,23 +1,29 @@ +"A class for testing input shapes on a PyTorch model." from typing import List, Optional, Dict -from model_inspector.utils import is_defined_shape, output_to_shape_dict +from src.model_inspector.utils import is_defined_shape, output_to_shape_dict import torch class InputTester: + """This class provides methods to test various input shapes + on a given PyTorch model and collect information + about working and non-working input sizes.""" def __init__(self, model, input_candidate=None): """ A class for testing input shapes on a PyTorch model. - This class provides methods to test various input shapes on a given PyTorch model and collect information + This class provides methods to test various input shapes + on a given PyTorch model and collect information about working and non-working input sizes. Args: model (torch.nn.Module): The PyTorch model to be tested. Note: - The try_image_input and try_audio_input methods test the model with predefined input sizes for image-like and - audio-like data, respectively. The get_shapes method retrieves the final input and output shapes after testing - various input sizes. + The try_image_input and try_audio_input methods test the + model with predefined input sizes for image-like and + audio-like data, respectively. The get_shapes method retrieves the final input + and output shapes after testing various input sizes. """ self.model = model self.input_candidate = input_candidate @@ -70,7 +76,6 @@ def try_image_input(self, n_dims: Optional[int] = None): """ rgb_size_1 = [3, 224, 224] rgb_size_2 = [3, 112, 112] - # TODO: add 'weirder' size like [3,113, 217] grey_size_1 = [1, 224, 224] grey_size_2 = [1, 112, 112] @@ -110,6 +115,7 @@ def try_audio_input(self): self.test_input_size(input_size) def test_input_size(self, input_size): + "Test a specific input size on the model." input_array = torch.ones( (input_size) ).numpy() # i get type issues when using np.ones() @@ -117,7 +123,7 @@ def test_input_size(self, input_size): output = None try: output = self.model.infer(input_tensor) - except Exception: + except Exception: #pylint: disable=(broad-exception-caught) pass if output is not None: self.working_input_sizes["input"].append(input_size) @@ -130,6 +136,8 @@ def test_input_size(self, input_size): self.working_output_sizes[output] = [shape] def try_inputs(self): + """Test candidate input (if provided), + image-like inputs, and audio-like inputs.""" if self.input_candidate: if is_defined_shape(self.input_candidate): self.test_input_size(self.input_candidate) @@ -137,6 +145,7 @@ def try_inputs(self): self.try_audio_input() def get_shapes(self): + "Perform tests on all inputs and retrieve the final input and output shapes." self.try_inputs() input_shapes, output_shapes = {}, {} for output_tensor_name, sizes in self.working_output_sizes.items(): diff --git a/src/model_inspector/inspector.py b/src/model_inspector/inspector.py index 032aacb..dfc8ebf 100644 --- a/src/model_inspector/inspector.py +++ b/src/model_inspector/inspector.py @@ -1,26 +1,45 @@ -import torch.nn as nn -import torch -from model_inspector.input_size_calculator import InputSizeCalculator -from model_inspector.input_tester import InputTester +""" +Module for inspecting and gathering metadata from a PyTorch model. +This module provides functionality to inspect a PyTorch model, +reverse its layers for input shape calculation, validate input shapes, +and retrieve metadata such as input and output tensor shapes. +""" from viam.services.mlmodel import Metadata, TensorInfo from viam.logging import getLogger from viam.utils import dict_to_struct +from src.model_inspector.input_size_calculator import InputSizeCalculator +from src.model_inspector.input_tester import InputTester + +from torch import nn +import torch + + LOGGER = getLogger(__name__) class Inspector: + "Inspector class for analyzing and gathering metadata from a PyTorch model." def __init__(self, model: nn.Module) -> None: # self.summary: ModelStatistics = summary(module, input_size=[1,3, 640,480]) self.model = model self.input_size_calculator = InputSizeCalculator() self.dimensionality = None + input_shape_candidate = self.reverse_module() + self.input_tester = InputTester(self.model, input_shape_candidate) def find_metadata(self, label_path: str): + """ + Gather metadata including input and output tensor information. + + Args: + label_path (str): Path to the label file, if available. + + Returns: + Metadata: Metadata object containing model information. + """ input_info, output_info = [], [] - input_shape_candidate = self.reverse_module() - self.input_tester = InputTester(self.model, input_shape_candidate) input_shapes, output_shapes = self.input_tester.get_shapes() for input_tensor_name, shape in input_shapes.items(): input_info.append(TensorInfo(name=input_tensor_name, shape=shape)) @@ -54,7 +73,6 @@ def reverse_module(self): modules = list(self.model.model.children()) modules.reverse() # last layer comes first so need to reverse it - # TODO: Add output shape from label files # if self.len_module_labels is not None: # output_shape = (1, self.len_module_labels) # else: @@ -67,13 +85,22 @@ def reverse_module(self): output_shape = self.input_size_calculator.get_input_size( module, input_shape ) - LOGGER.info(f"For module {module}, the output shape is {output_shape}") + LOGGER.info("For module %s, the output shape is %s", module, output_shape) else: continue # sometimes some children are None return input_shape def is_valid_input_shape(self, input_shape): + """ + Validate if a given input shape is valid for the model. + + Args: + input_shape: Shape of the input tensor to validate. + + Returns: + torch.Size or None: Size of the output tensor if input shape is valid, or None if not. + """ input_tensor = torch.ones(input_shape) try: output = self.model(input_tensor) diff --git a/src/model_inspector/test_input_size_calculator.py b/src/model_inspector/test_input_size_calculator.py deleted file mode 100644 index db1ba31..0000000 --- a/src/model_inspector/test_input_size_calculator.py +++ /dev/null @@ -1,97 +0,0 @@ -import torch.nn as nn -import torch -from inspector import Inspector -from torch.nn import functional as F -import unittest - - -class NotSequential1(nn.Module): - def __init__(self): - super(NotSequential1, self).__init__() - self.conv1 = nn.Conv2d(1, 20, 5) - self.conv2 = nn.Conv2d(20, 1, 5) - self.fully_connected = nn.Linear(70, 30) - - return super().__init_subclass__() - - def forward(self, x): - x = F.relu(self.conv1(x)) - x = F.relu(self.conv2(x)) - return F.softmax(self.fully_connected(x)) - - -class NotSequential2(nn.Module): - def __init__(self) -> None: - super(NotSequential2, self).__init__() - self.conv2 = nn.Conv2d(20, 1, 5) - self.conv1 = nn.Conv2d(1, 20, 5) - self.fully_connected = nn.Linear(70, 30) - - return super().__init_subclass__() - - def forward(self, x): - x = F.relu(self.conv1(x)) - x = F.relu(self.conv2(x)) - return F.softmax(self.fully_connected(x)) - - -class TestInputs(unittest.TestCase): - def __init__(self, methodName: str = "runTest") -> None: - super().__init__(methodName) - - self.sequential_model = nn.Sequential( - nn.Conv2d(1, 20, 5), - nn.ReLU(), - nn.Conv2d(20, 1, 5), - nn.ReLU(), - nn.Linear(70, 30), - ) - - self.not_sequential_model_1 = NotSequential1() - self.not_sequential_model_2 = NotSequential2() - self.inspector_seq = Inspector(self.not_sequential_model_1) - self.inspector_not_seq_1 = Inspector(self.not_sequential_model_1) - self.inspector_not_seq_2 = Inspector(self.not_sequential_model_2) - - self.valid_input_size = torch.Size([1, 9, 78]) - self.not_valid_input_size = torch.Size([1, 9, 9]) - - def test_inference_seq(self): - input_tensor = torch.ones(self.valid_input_size) - _ = self.sequential_model(input_tensor) - with self.assertRaises(RuntimeError): - input_tensor = torch.ones(self.not_valid_input_size) - _ = self.not_sequential_model_1(input_tensor) - - def test_inference_not_seq_1(self): - input_tensor = torch.ones(self.valid_input_size) - _ = self.not_sequential_model_1(input_tensor) - with self.assertRaises(RuntimeError): - input_tensor = torch.ones(self.not_valid_input_size) - _ = self.not_sequential_model_1(input_tensor) - - def test_inference_not_seq_2(self): - input_tensor = torch.ones(self.valid_input_size) - _ = self.not_sequential_model_2(input_tensor) - with self.assertRaises(RuntimeError): - input_tensor = torch.ones(self.not_valid_input_size) - _ = self.not_sequential_model_1(input_tensor) - - def test_found_input_shape_seq(self): - found_input_shape = torch.Size(self.inspector_seq.reverse_module()) - self.assertEqual(found_input_shape, self.valid_input_size) - self.assertNotEqual(found_input_shape, self.not_valid_input_size) - - def test_found_input_shape_not_seq_1(self): - found_input_shape = torch.Size(self.inspector_not_seq_1.reverse_module()) - self.assertEqual(found_input_shape, self.valid_input_size) - self.assertNotEqual(found_input_shape, self.not_valid_input_size) - - def test_found_input_shape_not_seq_2(self): - found_input_shape = torch.Size(self.inspector_not_seq_2.reverse_module()) - self.assertNotEqual(found_input_shape, self.valid_input_size) - self.assertNotEqual(found_input_shape, self.not_valid_input_size) - - -if __name__ == "__main__": - unittest.main() diff --git a/src/model_inspector/test_input_tester.py b/src/model_inspector/test_input_tester.py deleted file mode 100644 index 9e2ce4a..0000000 --- a/src/model_inspector/test_input_tester.py +++ /dev/null @@ -1,60 +0,0 @@ -import torch.nn as nn -import unittest -from input_tester import InputTester - - -class TestInputs(unittest.TestCase): - def __init__(self, methodName: str = "runTest") -> None: - super().__init__(methodName) - self.image_model = nn.Sequential( - nn.Conv2d(1, 20, 5), - nn.ReLU(), - nn.Conv2d(20, 1, 5), - nn.ReLU(), - ) - - self.nc = InputTester(self.image_model) - - def test_grey_image_model(self): - grey_image_model = nn.Sequential( - nn.Conv2d(1, 20, 5), - nn.ReLU(), - nn.Conv2d(20, 1, 5), - nn.ReLU(), - ) - input_tester = InputTester(grey_image_model) - input_shape = input_tester.get_shapes() - self.assertEqual(input_shape, [1, -1, -1]) - - def test_rgb_image_model(self): - grey_image_model = nn.Sequential( - nn.Conv2d(3, 20, 5), - nn.ReLU(), - nn.Conv2d(20, 1, 5), - nn.ReLU(), - ) - input_tester = InputTester(grey_image_model) - input_shape = input_tester.get_shapes() - self.assertEqual(input_shape, [3, -1, -1]) - - def test_mono_audio_model(self): - grey_image_model = nn.Sequential( - nn.Conv1d(1, 3, 5), - nn.ReLU(), - ) - input_tester = InputTester(grey_image_model) - input_shape = input_tester.get_shapes() - self.assertEqual(input_shape, [1, -1]) - - def test_stereo_audio_model(self): - grey_image_model = nn.Sequential( - nn.Conv1d(2, 3, 3), - nn.ReLU(), - ) - input_tester = InputTester(grey_image_model) - input_shape = input_tester.get_shapes() - self.assertEqual(input_shape, [2, -1]) - - -if __name__ == "__main__": - unittest.main() diff --git a/src/model_inspector/utils.py b/src/model_inspector/utils.py index 3ad43c8..fb9d5af 100644 --- a/src/model_inspector/utils.py +++ b/src/model_inspector/utils.py @@ -1,7 +1,10 @@ +"Utility functions for model inspection and input shape validation." + from typing import Tuple, Dict, List from numpy.typing import NDArray -import torch import numpy as np +import torch + def is_valid_input_shape(model, input_shape, add_batch_dimension: bool = False): @@ -10,12 +13,15 @@ def is_valid_input_shape(model, input_shape, add_batch_dimension: bool = False): Args: model (torch.nn.Module): The PyTorch model to validate the input shape for. - input_shape (tuple): The shape of the input tensor. It should be in the format (C, H, W) for image-like data, - where C is the number of channels, H is the height, and W is the width. - add_batch_dimension (bool, optional): Whether to add a batch dimension to the input tensor. Default is False. + input_shape (tuple): The shape of the input tensor. + It should be in the format (C, H, W) for image-like data, + where C is the number of channels, H is the height, and W is the width. + add_batch_dimension (bool, optional): + Whether to add a batch dimension to the input tensor. Default is False. Returns: - list or None: A list representing the shape of the output tensor if the input shape is valid for the model, + list or None: A list representing the shape of the output tensor + if the input shape is valid for the model, or None if an exception occurs during the model evaluation. """ @@ -29,7 +35,9 @@ def is_valid_input_shape(model, input_shape, add_batch_dimension: bool = False): except ValueError: return None if isinstance(output, torch.Tensor): - return list(output.size()) + out = list(output.size()) + return out + return None def is_defined_shape(shape: Tuple[int]) -> bool: diff --git a/src/torch_mlmodel_module.py b/src/torch_mlmodel_module.py index a4b802d..a50766c 100644 --- a/src/torch_mlmodel_module.py +++ b/src/torch_mlmodel_module.py @@ -1,3 +1,10 @@ +""" +This module provides functionality to infer input size predictions +and retrieve metadata associated with a model. + +The module initializes by loading a TorchModel from a specified model file path and +configures an Inspector to extract metadata, including labels if provided. +""" from typing import ClassVar, Mapping, Sequence, Dict, Optional from numpy.typing import NDArray from typing_extensions import Self @@ -9,28 +16,39 @@ from viam.resource.base import ResourceBase from viam.utils import ValueTypes from viam.logging import getLogger -from model.model import TorchModel -from model_inspector.inspector import Inspector +from src.model.model import TorchModel +from src.model_inspector.inspector import Inspector LOGGER = getLogger(__name__) class TorchMLModelModule(MLModel, Reconfigurable): + """ + This class integrates a PyTorch model with Viam's MLModel and Reconfigurable interfaces, + providing functionality to create, configure, and use the model for inference. + """ MODEL: ClassVar[Model] = Model(ModelFamily("viam", "mlmodel"), "torch-cpu") def __init__(self, name: str): super().__init__(name=name) + self.path_to_model_file = None + + self.torch_model = None + self.inspector = None + self._metadata = None @classmethod def new_service( cls, config: ServiceConfig, dependencies: Mapping[ResourceName, ResourceBase] ) -> Self: + "Create and configure a new instance of the service." service = cls(config.name) service.reconfigure(config, dependencies) return service @classmethod def validate_config(cls, config: ServiceConfig) -> Sequence[str]: + "Validate the configuration for the service." model_path = config.attributes.fields["model_path"].string_value if model_path == "": raise Exception( @@ -41,6 +59,8 @@ def validate_config(cls, config: ServiceConfig) -> Sequence[str]: def reconfigure( self, config: ServiceConfig, dependencies: Mapping[ResourceName, ResourceBase] ): + "Reconfigure the service with the given configuration and dependencies." + # pylint: disable=too-many-return-statements def get_attribute_from_config(attribute_name: str, default, of_type=None): if attribute_name not in config.attributes.fields: return default @@ -55,16 +75,15 @@ def get_attribute_from_config(attribute_name: str, default, of_type=None): type_default = type(default) if type_default == bool: return config.attributes.fields[attribute_name].bool_value - elif type_default == int: + if type_default == int: return int(config.attributes.fields[attribute_name].number_value) - elif type_default == float: + if type_default == float: return config.attributes.fields[attribute_name].number_value - elif type_default == str: + if type_default == str: return config.attributes.fields[attribute_name].string_value - elif type_default == dict: + if type_default == dict: return dict(config.attributes.fields[attribute_name].struct_value) - - # TODO: Test self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + return default self.path_to_model_file = get_attribute_from_config("model_path", None, str) label_file = get_attribute_from_config("label_path", None, str) @@ -73,21 +92,26 @@ def get_attribute_from_config(attribute_name: str, default, of_type=None): self.inspector = Inspector(self.torch_model) self._metadata = self.inspector.find_metadata(label_file) + async def infer( self, input_tensors: Dict[str, NDArray], *, timeout: Optional[float] ) -> Dict[str, NDArray]: - """Take an already ordered input tensor as an array, make an inference on the model, and return an output tensor map. + """Take an already ordered input tensor as an array, + make an inference on the model, and return an output tensor map. Args: - input_tensors (Dict[str, NDArray]): A dictionary of input flat tensors as specified in the metadata + input_tensors (Dict[str, NDArray]): + A dictionary of input flat tensors as specified in the metadata Returns: - Dict[str, NDArray]: A dictionary of output flat tensors as specified in the metadata + Dict[str, NDArray]: + A dictionary of output flat tensors as specified in the metadata """ return self.torch_model.infer(input_tensors) async def metadata(self, *, timeout: Optional[float]) -> Metadata: - """Get the metadata (such as name, type, expected tensor/array shape, inputs, and outputs) associated with the ML model. + """Get the metadata (such as name, type, expected tensor/array shape, + inputs, and outputs) associated with the ML model. Returns: Metadata: The metadata diff --git a/test_local.py b/tests/test_local.py similarity index 69% rename from test_local.py rename to tests/test_local.py index 0aca518..8156453 100644 --- a/test_local.py +++ b/tests/test_local.py @@ -1,17 +1,33 @@ +from google.protobuf.struct_pb2 import Struct import torch import unittest from src.model.model import TorchModel from src.model_inspector.inspector import Inspector +from src.torch_mlmodel_module import TorchMLModelModule from viam.services.mlmodel import Metadata +from viam.proto.app.robot import ComponentConfig from torchvision.models.detection import FasterRCNN from torchvision.models import MobileNet_V2_Weights import torchvision import os from torchvision.models.detection.rpn import AnchorGenerator +from typing import List, Iterable, Dict, Any, Mapping +import numpy as np -class TestInputs(unittest.TestCase): +def make_component_config(dictionary: Mapping[str, Any]) -> ComponentConfig: + struct = Struct() + struct.update(dictionary) + return ComponentConfig(attributes=struct) + +config = ( + make_component_config({"model_path": "model path"}), + "received only one dimension attribute", +) + + +class TestInputs(unittest.IsolatedAsyncioTestCase): @staticmethod def load_resnet_weights(): return TorchModel( @@ -53,6 +69,20 @@ def load_detector_from_torchvision(): def __init__(self, methodName: str = "runTest") -> None: super().__init__(methodName) + async def test_validate(self): + response = TorchMLModelModule.validate_config(config=config[0]) + self.assertEqual(response, []) + + async def test_validate_empty_config(self): + empty_config = make_component_config({}) + with self.assertRaises(Exception) as excinfo: + await TorchMLModelModule.validate_config(config=empty_config) + + self.assertIn( + "model_path can't be empty. model is required for torch mlmoded service module.", + str(excinfo.exception) + ) + def test_error_loading_weights(self): with self.assertRaises(TypeError): _ = self.load_resnet_weights() @@ -95,6 +125,17 @@ def test_detector_metadata(self): print(f"Checked {output_name} ") self.assertTrue(output_checked) + def test_infer_method(self): + model: TorchModel = self.load_detector_from_torchvision() + x = torch.ones(3, 300, 400).unsqueeze(0) + output = model.infer({"input_name": x.numpy()}) + self.assertIsInstance(output, dict) + + # Assert the structure of the output based on wrap_output function + for key, value in output.items(): + self.assertIsInstance(key, str) + self.assertIsInstance(value, np.ndarray) + if __name__ == "__main__": unittest.main() From 939cdb6d455c7f6e0c76ce30866d9b58ba580948 Mon Sep 17 00:00:00 2001 From: dhritinaidu Date: Fri, 12 Jul 2024 12:33:45 -0400 Subject: [PATCH 03/11] changed action version in publish.yaml --- .github/workflows/publish.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index 1235597..695af3f 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -8,7 +8,7 @@ jobs: publish: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v3 + - uses: actions/checkout@v4 - uses: viamrobotics/build-action@v4 with: # note: you can replace this line with 'version: ""' if you want to test the build process without deploying From 7d7d38dc03a1426863063363c8a15e7692f57e0f Mon Sep 17 00:00:00 2001 From: dhritinaidu Date: Fri, 12 Jul 2024 12:37:28 -0400 Subject: [PATCH 04/11] updated meta.json --- meta.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/meta.json b/meta.json index 1d681c6..f2e855d 100644 --- a/meta.json +++ b/meta.json @@ -12,7 +12,7 @@ "build": { "build": "./build.sh", "path": "dist/archive.tar.gz", - "arch": ["linux/arm64", "linux/amd64"] + "arch": ["darwin/arm64", "darwin/amd64", "linux/arm64", "linux/amd64"] }, "entrypoint": "dist/main" } \ No newline at end of file From b64e7a5fb8dcc49787c98e7ffad74d790e21762f Mon Sep 17 00:00:00 2001 From: dhritinaidu Date: Fri, 12 Jul 2024 12:14:19 -0400 Subject: [PATCH 05/11] added workflows and unit tests fixed merge conflicts --- meta.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/meta.json b/meta.json index f2e855d..1d681c6 100644 --- a/meta.json +++ b/meta.json @@ -12,7 +12,7 @@ "build": { "build": "./build.sh", "path": "dist/archive.tar.gz", - "arch": ["darwin/arm64", "darwin/amd64", "linux/arm64", "linux/amd64"] + "arch": ["linux/arm64", "linux/amd64"] }, "entrypoint": "dist/main" } \ No newline at end of file From d9131995f5585121e0dc2579f3c69d4f1b1fd8be Mon Sep 17 00:00:00 2001 From: dhritinaidu Date: Fri, 12 Jul 2024 15:35:23 -0400 Subject: [PATCH 06/11] updated requirements --- requirements.txt | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index 23e856d..af2b6e6 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,7 @@ viam-sdk typing-extensions -numpy==1.24.1 +numpy<2.0.0 +pylint pyinstaller google-api-python-client torchvision From 298b1903d779a8e3b015ca6072b7d3d930599679 Mon Sep 17 00:00:00 2001 From: dhritinaidu Date: Fri, 12 Jul 2024 16:15:32 -0400 Subject: [PATCH 07/11] fixed linting issues --- Makefile | 4 +- src/main.py | 1 - src/model/model.py | 5 +- src/model_inspector/input_size_calculator.py | 4 +- src/model_inspector/input_tester.py | 19 +++--- src/model_inspector/inspector.py | 9 ++- src/model_inspector/utils.py | 7 +-- {tests => src}/test_local.py | 65 ++++++++++++++++---- src/torch_mlmodel_module.py | 15 ++--- 9 files changed, 86 insertions(+), 43 deletions(-) rename {tests => src}/test_local.py (78%) diff --git a/Makefile b/Makefile index 31f9b93..b0e6834 100644 --- a/Makefile +++ b/Makefile @@ -3,9 +3,9 @@ PYTHONPATH := ./torch:$(PYTHONPATH) .PHONY: test lint setup dist test: - PYTHONPATH=$(PYTHONPATH) pytest tests/ + PYTHONPATH=$(PYTHONPATH) pytest src/ lint: - pylint --disable=E1101,W0719,C0202,R0801,W0613 src/ + pylint --disable=E1101,W0719,C0202,R0801,W0613,C0411 src/ setup: python3 -m pip install -r requirements.txt -U dist/archive.tar.gz: diff --git a/src/main.py b/src/main.py index c1cca5e..e22981f 100644 --- a/src/main.py +++ b/src/main.py @@ -30,5 +30,4 @@ async def main(): if __name__ == "__main__": - asyncio.run(main()) diff --git a/src/model/model.py b/src/model/model.py index 69fb950..f550b57 100644 --- a/src/model/model.py +++ b/src/model/model.py @@ -36,8 +36,7 @@ def __init__( if size_mb > 500: # pylint: disable=deprecated-method LOGGER.warn( - "model file may be large for certain hardware (%s MB)", - size_mb + "model file may be large for certain hardware (%s MB)", size_mb ) self.model = torch.load(path_to_serialized_file) if not isinstance(self.model, nn.Module): @@ -47,7 +46,7 @@ def __init__( is of type collections.OrderedDict, which suggests that the provided file describes weights instead of a standalone model""", - path_to_serialized_file + path_to_serialized_file, ) raise TypeError( f"the model is of type {type(self.model)} instead of nn.Module type" diff --git a/src/model_inspector/input_size_calculator.py b/src/model_inspector/input_size_calculator.py index e094591..4ad7147 100644 --- a/src/model_inspector/input_size_calculator.py +++ b/src/model_inspector/input_size_calculator.py @@ -10,7 +10,7 @@ from typing import Dict, Tuple from viam.logging import getLogger from torch import nn -from src.model_inspector.utils import is_defined_shape +from model_inspector.utils import is_defined_shape LOGGER = getLogger(__name__) @@ -312,6 +312,6 @@ def get_input_size( "For layer type %s, with output shape: %s input shape found is %s", type(layer).__name__, output_shape, - input_shape + input_shape, ) return input_shape diff --git a/src/model_inspector/input_tester.py b/src/model_inspector/input_tester.py index 051cbab..0995c3f 100644 --- a/src/model_inspector/input_tester.py +++ b/src/model_inspector/input_tester.py @@ -1,18 +1,19 @@ "A class for testing input shapes on a PyTorch model." from typing import List, Optional, Dict -from src.model_inspector.utils import is_defined_shape, output_to_shape_dict +from model_inspector.utils import is_defined_shape, output_to_shape_dict import torch class InputTester: - """This class provides methods to test various input shapes - on a given PyTorch model and collect information - about working and non-working input sizes.""" + """This class provides methods to test various input shapes + on a given PyTorch model and collect information + about working and non-working input sizes.""" + def __init__(self, model, input_candidate=None): """ A class for testing input shapes on a PyTorch model. - This class provides methods to test various input shapes + This class provides methods to test various input shapes on a given PyTorch model and collect information about working and non-working input sizes. @@ -20,9 +21,9 @@ def __init__(self, model, input_candidate=None): model (torch.nn.Module): The PyTorch model to be tested. Note: - The try_image_input and try_audio_input methods test the + The try_image_input and try_audio_input methods test the model with predefined input sizes for image-like and - audio-like data, respectively. The get_shapes method retrieves the final input + audio-like data, respectively. The get_shapes method retrieves the final input and output shapes after testing various input sizes. """ self.model = model @@ -123,7 +124,7 @@ def test_input_size(self, input_size): output = None try: output = self.model.infer(input_tensor) - except Exception: #pylint: disable=(broad-exception-caught) + except Exception: # pylint: disable=(broad-exception-caught) pass if output is not None: self.working_input_sizes["input"].append(input_size) @@ -136,7 +137,7 @@ def test_input_size(self, input_size): self.working_output_sizes[output] = [shape] def try_inputs(self): - """Test candidate input (if provided), + """Test candidate input (if provided), image-like inputs, and audio-like inputs.""" if self.input_candidate: if is_defined_shape(self.input_candidate): diff --git a/src/model_inspector/inspector.py b/src/model_inspector/inspector.py index dfc8ebf..f1a2fa6 100644 --- a/src/model_inspector/inspector.py +++ b/src/model_inspector/inspector.py @@ -8,8 +8,8 @@ from viam.logging import getLogger from viam.utils import dict_to_struct -from src.model_inspector.input_size_calculator import InputSizeCalculator -from src.model_inspector.input_tester import InputTester +from model_inspector.input_size_calculator import InputSizeCalculator +from model_inspector.input_tester import InputTester from torch import nn import torch @@ -20,6 +20,7 @@ class Inspector: "Inspector class for analyzing and gathering metadata from a PyTorch model." + def __init__(self, model: nn.Module) -> None: # self.summary: ModelStatistics = summary(module, input_size=[1,3, 640,480]) self.model = model @@ -85,7 +86,9 @@ def reverse_module(self): output_shape = self.input_size_calculator.get_input_size( module, input_shape ) - LOGGER.info("For module %s, the output shape is %s", module, output_shape) + LOGGER.info( + "For module %s, the output shape is %s", module, output_shape + ) else: continue # sometimes some children are None diff --git a/src/model_inspector/utils.py b/src/model_inspector/utils.py index fb9d5af..b8bf56c 100644 --- a/src/model_inspector/utils.py +++ b/src/model_inspector/utils.py @@ -6,21 +6,20 @@ import torch - def is_valid_input_shape(model, input_shape, add_batch_dimension: bool = False): """ Check if the input shape is valid for a given PyTorch model. Args: model (torch.nn.Module): The PyTorch model to validate the input shape for. - input_shape (tuple): The shape of the input tensor. + input_shape (tuple): The shape of the input tensor. It should be in the format (C, H, W) for image-like data, where C is the number of channels, H is the height, and W is the width. - add_batch_dimension (bool, optional): + add_batch_dimension (bool, optional): Whether to add a batch dimension to the input tensor. Default is False. Returns: - list or None: A list representing the shape of the output tensor + list or None: A list representing the shape of the output tensor if the input shape is valid for the model, or None if an exception occurs during the model evaluation. """ diff --git a/tests/test_local.py b/src/test_local.py similarity index 78% rename from tests/test_local.py rename to src/test_local.py index 8156453..bdf7bb1 100644 --- a/tests/test_local.py +++ b/src/test_local.py @@ -1,26 +1,37 @@ -from google.protobuf.struct_pb2 import Struct -import torch +"Module for unit testing functionalities related to TorchModel and TorchMLModelModule." + +from google.protobuf.struct_pb2 import Struct #pylint: disable=(no-name-in-module) + import unittest -from src.model.model import TorchModel -from src.model_inspector.inspector import Inspector -from src.torch_mlmodel_module import TorchMLModelModule -from viam.services.mlmodel import Metadata -from viam.proto.app.robot import ComponentConfig +from torchvision.models.detection.rpn import AnchorGenerator + from torchvision.models.detection import FasterRCNN from torchvision.models import MobileNet_V2_Weights import torchvision -import os -from torchvision.models.detection.rpn import AnchorGenerator -from typing import List, Iterable, Dict, Any, Mapping +from model.model import TorchModel +from model_inspector.inspector import Inspector +from torch_mlmodel_module import TorchMLModelModule +from viam.services.mlmodel import Metadata +from viam.proto.app.robot import ComponentConfig + +from typing import Any, Mapping import numpy as np +import os + + + +import torch + def make_component_config(dictionary: Mapping[str, Any]) -> ComponentConfig: + " makes a mock config" struct = Struct() struct.update(dictionary) return ComponentConfig(attributes=struct) + config = ( make_component_config({"model_path": "model path"}), "received only one dimension attribute", @@ -28,8 +39,14 @@ def make_component_config(dictionary: Mapping[str, Any]) -> ComponentConfig: class TestInputs(unittest.IsolatedAsyncioTestCase): + """ + Unit tests for validating TorchModel and TorchMLModelModule functionalities. + """ @staticmethod def load_resnet_weights(): + """ + Load ResNet weights from a serialized file. + """ return TorchModel( path_to_serialized_file=os.path.join( "examples", "resnet_18", "resnet18-f37072fd.pth" @@ -38,6 +55,9 @@ def load_resnet_weights(): @staticmethod def load_standalone_resnet(): + """ + Load a standalone ResNet model. + """ return TorchModel( path_to_serialized_file=os.path.join( "examples", "resnet_18_scripted", "resnet-18.pt" @@ -46,6 +66,9 @@ def load_standalone_resnet(): @staticmethod def load_detector_from_torchvision(): + """ + Load a detector model using torchvision. + """ backbone = torchvision.models.mobilenet_v2( weights=MobileNet_V2_Weights.DEFAULT ).features @@ -66,28 +89,40 @@ def load_detector_from_torchvision(): model.eval() return TorchModel(path_to_serialized_file=None, model=model) - def __init__(self, methodName: str = "runTest") -> None: + def __init__(self, methodName: str = "runTest") -> None: #pylint: disable=(useless-parent-delegation) super().__init__(methodName) async def test_validate(self): + """ + Test validation of configuration using TorchMLModelModule. + """ response = TorchMLModelModule.validate_config(config=config[0]) self.assertEqual(response, []) async def test_validate_empty_config(self): + """ + Test validation with an empty configuration. + """ empty_config = make_component_config({}) with self.assertRaises(Exception) as excinfo: await TorchMLModelModule.validate_config(config=empty_config) self.assertIn( "model_path can't be empty. model is required for torch mlmoded service module.", - str(excinfo.exception) + str(excinfo.exception), ) def test_error_loading_weights(self): + """ + Test error handling when loading ResNet weights. + """ with self.assertRaises(TypeError): _ = self.load_resnet_weights() def test_resnet_metadata(self): + """ + Test metadata retrieval for ResNet model. + """ model: TorchModel = self.load_standalone_resnet() x = torch.ones(3, 300, 400).unsqueeze(0) output = model.infer({"any_input_name_you_want": x.numpy()}) @@ -107,6 +142,9 @@ def test_resnet_metadata(self): self.assertTrue(output_checked) def test_detector_metadata(self): + """ + Test metadata retrieval for detector model. + """ model: TorchModel = self.load_detector_from_torchvision() x = torch.ones(3, 300, 400).unsqueeze(0) output = model.infer({"any_input_name_you_want": x.numpy()}) @@ -126,6 +164,9 @@ def test_detector_metadata(self): self.assertTrue(output_checked) def test_infer_method(self): + """ + Test inference method of the detector model. + """ model: TorchModel = self.load_detector_from_torchvision() x = torch.ones(3, 300, 400).unsqueeze(0) output = model.infer({"input_name": x.numpy()}) diff --git a/src/torch_mlmodel_module.py b/src/torch_mlmodel_module.py index a50766c..ff8b90b 100644 --- a/src/torch_mlmodel_module.py +++ b/src/torch_mlmodel_module.py @@ -16,8 +16,8 @@ from viam.resource.base import ResourceBase from viam.utils import ValueTypes from viam.logging import getLogger -from src.model.model import TorchModel -from src.model_inspector.inspector import Inspector +from model.model import TorchModel +from model_inspector.inspector import Inspector LOGGER = getLogger(__name__) @@ -27,6 +27,7 @@ class TorchMLModelModule(MLModel, Reconfigurable): This class integrates a PyTorch model with Viam's MLModel and Reconfigurable interfaces, providing functionality to create, configure, and use the model for inference. """ + MODEL: ClassVar[Model] = Model(ModelFamily("viam", "mlmodel"), "torch-cpu") def __init__(self, name: str): @@ -60,6 +61,7 @@ def reconfigure( self, config: ServiceConfig, dependencies: Mapping[ResourceName, ResourceBase] ): "Reconfigure the service with the given configuration and dependencies." + # pylint: disable=too-many-return-statements def get_attribute_from_config(attribute_name: str, default, of_type=None): if attribute_name not in config.attributes.fields: @@ -92,25 +94,24 @@ def get_attribute_from_config(attribute_name: str, default, of_type=None): self.inspector = Inspector(self.torch_model) self._metadata = self.inspector.find_metadata(label_file) - async def infer( self, input_tensors: Dict[str, NDArray], *, timeout: Optional[float] ) -> Dict[str, NDArray]: - """Take an already ordered input tensor as an array, + """Take an already ordered input tensor as an array, make an inference on the model, and return an output tensor map. Args: - input_tensors (Dict[str, NDArray]): + input_tensors (Dict[str, NDArray]): A dictionary of input flat tensors as specified in the metadata Returns: - Dict[str, NDArray]: + Dict[str, NDArray]: A dictionary of output flat tensors as specified in the metadata """ return self.torch_model.infer(input_tensors) async def metadata(self, *, timeout: Optional[float]) -> Metadata: - """Get the metadata (such as name, type, expected tensor/array shape, + """Get the metadata (such as name, type, expected tensor/array shape, inputs, and outputs) associated with the ML model. Returns: From 27d02aa6ee6b848b5ed5c5e74d3266ea4c0e432d Mon Sep 17 00:00:00 2001 From: dhritinaidu Date: Fri, 12 Jul 2024 16:25:10 -0400 Subject: [PATCH 08/11] added pytest as a requirement --- requirements.txt | 1 + src/test_local.py | 2 ++ 2 files changed, 3 insertions(+) diff --git a/requirements.txt b/requirements.txt index af2b6e6..2cc6d29 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,5 +4,6 @@ numpy<2.0.0 pylint pyinstaller google-api-python-client +pytest torchvision torch==2.2.1 \ No newline at end of file diff --git a/src/test_local.py b/src/test_local.py index bdf7bb1..9187a1f 100644 --- a/src/test_local.py +++ b/src/test_local.py @@ -180,3 +180,5 @@ def test_infer_method(self): if __name__ == "__main__": unittest.main() + + From a3a293c1093fe8d95f35a6ff27ad52416482165c Mon Sep 17 00:00:00 2001 From: dhritinaidu Date: Fri, 12 Jul 2024 16:30:49 -0400 Subject: [PATCH 09/11] fixed more linting --- src/test_local.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/test_local.py b/src/test_local.py index 9187a1f..85b286b 100644 --- a/src/test_local.py +++ b/src/test_local.py @@ -180,5 +180,4 @@ def test_infer_method(self): if __name__ == "__main__": unittest.main() - - + \ No newline at end of file From 646c1bc006ea757b408a02803458415e1f0c40b8 Mon Sep 17 00:00:00 2001 From: dhritinaidu Date: Fri, 12 Jul 2024 16:39:36 -0400 Subject: [PATCH 10/11] updated makefile for pyinstaller --- Makefile | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/Makefile b/Makefile index b0e6834..f5e1a71 100644 --- a/Makefile +++ b/Makefile @@ -1,12 +1,12 @@ PYTHONPATH := ./torch:$(PYTHONPATH) -.PHONY: test lint setup dist +.PHONY: test lint build dist test: PYTHONPATH=$(PYTHONPATH) pytest src/ lint: pylint --disable=E1101,W0719,C0202,R0801,W0613,C0411 src/ -setup: - python3 -m pip install -r requirements.txt -U +build: + ./build.sh dist/archive.tar.gz: - tar -czf module.tar.gz run.sh requirements.txt src \ No newline at end of file + tar -czvf dist/archive.tar.gz dist/__main__ \ No newline at end of file From c5180ee3c2c83c639534daf30a78fa0896a50747 Mon Sep 17 00:00:00 2001 From: dhritinaidu Date: Fri, 12 Jul 2024 17:20:45 -0400 Subject: [PATCH 11/11] fixed the make file entrypoint --- Makefile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Makefile b/Makefile index f5e1a71..898513e 100644 --- a/Makefile +++ b/Makefile @@ -9,4 +9,4 @@ lint: build: ./build.sh dist/archive.tar.gz: - tar -czvf dist/archive.tar.gz dist/__main__ \ No newline at end of file + tar -czvf dist/archive.tar.gz dist/main \ No newline at end of file