diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml new file mode 100644 index 0000000..695af3f --- /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@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 + 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 index 0a18bc8..cfb2286 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -1,4 +1,5 @@ -name: Run lint +name: Run lint & unit tests + on: push: @@ -8,7 +9,8 @@ on: jobs: build: - name: "Run lint" + + name: "Run unit tests" runs-on: ubuntu-latest steps: @@ -20,3 +22,7 @@ jobs: - name: Run lint run: | make lint + + - name: Run unit tests + run: make test + 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..898513e --- /dev/null +++ b/Makefile @@ -0,0 +1,12 @@ +PYTHONPATH := ./torch:$(PYTHONPATH) + +.PHONY: test lint build dist + +test: + PYTHONPATH=$(PYTHONPATH) pytest src/ +lint: + pylint --disable=E1101,W0719,C0202,R0801,W0613,C0411 src/ +build: + ./build.sh +dist/archive.tar.gz: + tar -czvf dist/archive.tar.gz dist/main \ 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 2f13312..2cc6d29 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,9 @@ viam-sdk -numpy +typing-extensions +numpy<2.0.0 +pylint pyinstaller google-api-python-client -torch==2.2.1 +pytest +torchvision +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..e22981f 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(): """ diff --git a/src/model/model.py b/src/model/model.py index c5ae9af..f550b57 100644 --- a/src/model/model.py +++ b/src/model/model.py @@ -1,53 +1,73 @@ -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 +75,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 +90,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..4ad7147 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 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..0995c3f 100644 --- a/src/model_inspector/input_tester.py +++ b/src/model_inspector/input_tester.py @@ -1,23 +1,30 @@ +"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 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 +77,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 +116,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 +124,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 +137,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 +146,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..f1a2fa6 100644 --- a/src/model_inspector/inspector.py +++ b/src/model_inspector/inspector.py @@ -1,26 +1,46 @@ -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 model_inspector.input_size_calculator import InputSizeCalculator +from 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 +74,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 +86,24 @@ 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..b8bf56c 100644 --- a/src/model_inspector/utils.py +++ b/src/model_inspector/utils.py @@ -1,7 +1,9 @@ +"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 +12,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 +34,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/test_local.py b/src/test_local.py similarity index 56% rename from test_local.py rename to src/test_local.py index 0aca518..85b286b 100644 --- a/test_local.py +++ b/src/test_local.py @@ -1,19 +1,52 @@ -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 viam.services.mlmodel import Metadata +from torchvision.models.detection.rpn import AnchorGenerator + from torchvision.models.detection import FasterRCNN from torchvision.models import MobileNet_V2_Weights import torchvision +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 -from torchvision.models.detection.rpn import AnchorGenerator -class TestInputs(unittest.TestCase): + +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", +) + + +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" @@ -22,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" @@ -30,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 @@ -50,14 +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), + ) + 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()}) @@ -77,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()}) @@ -95,6 +163,21 @@ def test_detector_metadata(self): print(f"Checked {output_name} ") 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()}) + 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() + \ No newline at end of file diff --git a/src/torch_mlmodel_module.py b/src/torch_mlmodel_module.py index a4b802d..ff8b90b 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 @@ -16,21 +23,33 @@ 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 +60,9 @@ 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 +77,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) @@ -76,18 +97,22 @@ def get_attribute_from_config(attribute_name: str, default, of_type=None): 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