Add a simple setup scripts (just for development)

nn_meter/ir_converters/utils.py

@@ -6,6 +6,7 @@ import json
 from .onnx_converter import OnnxConverter
 from .frozenpb_converter import FrozenPbConverter
 from .torch_converter import TorchConverter
+from .torch_converter.converter import NNIIRConverter
 
 
 def model_to_graph(model, model_type, input_shape=(1, 3, 224, 224)):
@@ -25,6 +26,9 @@ def model_to_graph(model, model_type, input_shape=(1, 3, 224, 224)):
             args = args.to("cuda")
         converter = TorchConverter(model, args)
         result = converter.convert()
+    elif model_type == 'nni':
+        converter = NNIIRConverter(model)
+        result = converter.convert()
     else:
         raise ValueError(f"Unsupported model type: {model_type}")
 

setup.py

@@ -0,0 +1,12 @@
+from setuptools import setup, find_packages
+
+
+setup(
+    name='nn_meter',
+    version='1.0',
+    description='',
+    author='',
+    author_email='',
+    url='',
+    packages=find_packages(),
+)