Dataset generator release

.github/workflows/integration-test_onnx_based_torch.yml

@@ -28,14 +28,14 @@ jobs:
         path: |
           ~/.nn_meter
           /home/runner/work/nn-Meter/data/testmodels
-        key: ${{hashFiles('nn_meter/configs/predictors.yaml')}}-${{hashFiles('tests/integration_test/test_latency_predictor.py')}}
+        key: ${{hashFiles('nn_meter/configs/predictors.yaml')}}--${{hashFiles('tests/integration_test/test_latency_predictor.py')}}
 
     - name: Install dependencies
       run: |
         pip install onnx==1.9.0
         pip install torch==1.9.0
         pip install torchvision==0.10.0
-        pip install onnx-simplifier
+        pip install onnx-simplifier==0.3.6
 
     - name: Install nn-Meter
       run: pip install -U .

README.md

@@ -48,7 +48,7 @@ nn-Meter is a latency predictor of models with type of Tensorflow, PyTorch, Onnx
 | Testing Model Type |                                                       Requirements                                                       |
 | :----------------: | :-----------------------------------------------------------------------------------------------------------------------: |
 |     Tensorflow     |                                                  `tensorflow==2.6.0`                                                  |
-|       Torch       | `torch==1.9.0`, `torchvision==0.10.0`, (alternative)[`onnx==1.9.0`, `onnx-simplifier==0.3.6`] or [`nni>=2.4`][1] |
+|       Torch       | `torch==1.9.0`, `torchvision==0.10.0`, (alternative)[`onnx>=1.9.0`, `onnx-simplifier==0.3.6`] or [`nni>=2.4`][1] |
 |        Onnx        |                                                      `onnx==1.9.0`                                                      |
 | nn-Meter IR graph |                                                            ---                                                            |
 |    NNI IR graph    |                                                       `nni>=2.4`                                                       |

examples/block-wise_latency_demo/test.py

@@ -0,0 +1,9 @@
+import pickle
+with open("/data/jiahang/working/nnmeterdata/predictor/adreno630gpu_tflite21/conv-bn-relu.pkl", "rb") as f:
+    predictor = pickle.load(f)
+
+print(predictor.predict([[1,2,3,4]]))
+print(predictor.predict([[1,2,3,]]))
+print(predictor.predict([[1,2]]))
+print(predictor.predict([[1]]))
+print(1)

examples/test_dataset_generator/test_generator.py

@@ -0,0 +1,3 @@
+from nn_meter.dataset.generator.generate_model import ModelGenerator
+generator = ModelGenerator("/data/jiahang/working/nn-Meter/nn_meter/dataset/generator/configs/vgg.yaml", "/data/jiahang/working/nn-Meter/examples/test_dataset_generator/")
+generator.run()

examples/test_dataset_generator/test_nasbench201.py

@@ -0,0 +1,13 @@
+config = "|nor_conv_3x3~0|+|nor_conv_3x3~0|nor_conv_1x1~1|+|skip_connect~0|nor_conv_3x3~1|nor_conv_3x3~2|"
+output = "/data1/jiahang/working/pixel4_int8_workspace/code/test.h5"
+from nn_meter.dataset.generator.networks.tf_network.nasbench201 import nasbench201
+nasbench201(config, output)
+
+
+
+
+from nas_201_api import NASBench201API as API
+api = API("/data/jiahang/working/NAS-Bench-201/NAS-Bench-201-v1_0-e61699.pth")
+index = api.query_index_by_arch(config)
+print(index)
+api.show(index)

nn_meter/builder/nn_modules/tf_networks/utils.py

@@ -4,14 +4,22 @@
 from tensorflow import keras
 
 def get_tensor_by_shapes(shapes, batch_size = 1):
+    """Generate tensor according to given shape
+    """
     if len(shapes) == 1:
         return tf.random.normal(shape = [batch_size] + shapes[0])
     else:
         return [tf.random.normal(shape = [batch_size] + shape) for shape in shapes]
 
 
 def get_inputs_by_shapes(shapes, batch_size = 1):
+    """ generate input tensor according to given shape
+    """
     if len(shapes) == 1:
         return keras.Input(shape=shapes[0], batch_size=batch_size)
     else:
         return [keras.Input(shape=shape, batch_size=batch_size) for shape in shapes]
+
+
+def save_model(model, save_path):
+    keras.models.save_model(model, save_path)

nn_meter/dataset/generator/configs/alexnet.yaml

@@ -0,0 +1,12 @@
+model_family: alexnet
+sample_count: 2000
+checkpoint_dir: ./data
+submodule_level: 2
+n_classes: 1000
+input_shape: [3,224,224]
+sample_space:
+    channel:
+       start: 0.2
+       end: 1.8
+       step: 0.1
+    kernelsize: [1,3,5,7]

nn_meter/dataset/generator/configs/densenet.yaml

@@ -0,0 +1,16 @@
+model_family: densenet
+sample_count: 2000
+checkpoint_dir: ./data
+submodule_level: 2
+n_classes: 1000
+input_shape: [3, 224, 224]
+sample_space:
+    channel:
+       start: 0.2
+       end: 1.2
+       step: 0.1
+    kernelsize: [1, 3, 5]
+    es:
+       start: 0.8
+       end: 1.8
+       step: 0.1

nn_meter/dataset/generator/configs/googlenet.yaml

@@ -0,0 +1,16 @@
+model_family: googlenet
+sample_count: 2000
+checkpoint_dir: ./data
+submodule_level: 2
+n_classes: 1000
+input_shape: [3, 224, 224]
+sample_space:
+    channel:
+       start: 0.2
+       end: 1.8
+       step: 0.1
+    kernelsize: [1, 3, 5, 7]
+    es:
+       start: 0.8
+       end: 1.8
+       step: 0.1

nn_meter/dataset/generator/configs/mnasnet.yaml

@@ -0,0 +1,13 @@
+model_family: mnasnet
+sample_count: 2000
+checkpoint_dir: ./data
+submodule_level: 2
+n_classes: 1000
+input_shape: [3, 224, 224]
+sample_space:
+    channel:
+       start: 0.2
+       end: 1.8
+       step: 0.1
+    kernelsize: [3, 5, 7]
+    es: [3, 4, 6]

nn_meter/dataset/generator/configs/mobilenetv1.yaml

@@ -0,0 +1,12 @@
+model_family: mobilenetv1
+sample_count: 2000
+checkpoint_dir: ./data
+submodule_level: 2
+n_classes: 1000
+input_shape: [3, 224, 224]
+sample_space:
+    channel:
+       start: 0.2
+       end: 1.8
+       step: 0.1
+    kernelsize: [1, 3, 5, 7]

nn_meter/dataset/generator/configs/mobilenetv2.yaml

@@ -0,0 +1,13 @@
+model_family: mobilenetv2
+sample_count: 2000
+checkpoint_dir: ./data
+submodule_level: 2
+n_classes: 1000
+input_shape: [3, 224, 224]
+sample_space:
+    channel:
+       start: 0.2
+       end: 1.8
+       step: 0.1
+    kernelsize: [1, 3, 5, 7]
+    es: [1, 2, 3, 4, 5, 6, 7, 8]

nn_meter/dataset/generator/configs/mobilenetv3.yaml

@@ -0,0 +1,17 @@
+model_family: mobilenetv3
+modelids: ['large','small']
+sample_count: 1000
+checkpoint_dir: ./data
+submodule_level: 2
+n_classes: 1000
+input_shape: [3, 224, 224]
+sample_space:
+    channel:
+       start: 0.2
+       end: 1.8
+       step: 0.1
+    kernelsize: [1, 3, 5, 7]
+    es:
+       start: 0.8
+       end: 1.8
+       step: 0.1

nn_meter/dataset/generator/configs/proxylessnas.yaml

@@ -0,0 +1,13 @@
+model_family: proxylessnas
+sample_count: 2000
+checkpoint_dir: ./data
+submodule_level: 2
+n_classes: 1000
+input_shape: [3, 224, 224]
+sample_space:
+    channel:
+       start: 0.2
+       end: 1.8
+       step: 0.1
+    kernelsize: [3, 5, 7]
+    es: [3, 6]

nn_meter/dataset/generator/configs/resnet.yaml

@@ -0,0 +1,14 @@
+model_family: resnet
+modelids: [18, 34, 50]
+sample_count: 500
+checkpoint_dir: ./data
+submodule_level: 2
+n_classes: 1000
+input_shape: [3, 224, 224]
+sample_space:
+    channel:
+       start: 0.2
+       end: 1.2
+       step: 0.1
+    kernelsize: [3, 5, 7]
+    es: [2, 4]

nn_meter/dataset/generator/configs/shufflenetv2.yaml

@@ -0,0 +1,16 @@
+model_family: shufflenetv2
+sample_count: 2000
+checkpoint_dir: ./data
+submodule_level: 2
+n_classes: 1000
+input_shape: [3, 224, 224]
+sample_space:
+    mid_channel:
+       start: 0.2
+       end: 1.8
+       step: 0.1
+    channel:
+       start: 1
+       end: 1.8
+       step: 0.1
+    kernelsize: [3, 5, 7, 9]

nn_meter/dataset/generator/configs/squeezenet.yaml

@@ -0,0 +1,12 @@
+model_family: squeezenet
+sample_count: 2000
+checkpoint_dir: ./data
+submodule_level: 2
+n_classes: 1000
+input_shape: [3, 224, 224]
+sample_space:
+    channel:
+       start: 0.2
+       end: 1.8
+       step: 0.1 
+    kernelsize: [1, 3, 5, 7]

nn_meter/dataset/generator/configs/vgg.yaml

@@ -0,0 +1,14 @@
+model_family: vgg
+modelids: [11, 13, 16, 19]
+sample_count: 500
+checkpoint_dir: ./data
+submodule_level: 2
+n_classes: 1000
+input_shape: [3, 224, 224]
+sample_space:
+    channel:
+       start: 0.2
+       end: 1.8
+       step: 0.1
+    kernelsize: [1, 3, 5, 7]
+

nn_meter/dataset/generator/generate_model.py

@@ -0,0 +1,83 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+import os
+import yaml
+import json
+import random
+from tensorflow import keras
+
+from .networks.tf_network import *
+
+
+MODELZOO = {
+    'alexnet':lambda input_tensor, cfg, version, flag: AlexNet(input_tensor, cfg, version, flag),
+    'mobilenetv1':lambda input_tensor, cfg, version, flag: MobileNetV1(input_tensor, cfg, version, flag),
+    'mobilenetv2':lambda input_tensor, cfg, version, flag: MobileNetV2(input_tensor, cfg, version, flag),
+    'mobilenetv3':lambda input_tensor, cfg, version, flag: MobileNetV3(input_tensor, cfg, version, flag),
+    'shufflenetv2':lambda input_tensor, cfg, version, flag: ShuffleNetV2(input_tensor, cfg, version, flag),
+    'vgg':lambda input_tensor, cfg, version, flag: VGG(input_tensor, cfg, version, flag), # vgg11/ 13/ 16/ 19
+    'resnet':lambda input_tensor, cfg, version, flag: ResNetV1(input_tensor, cfg, version, flag),
+    'proxylessnas':lambda input_tensor, cfg, version, flag: ProxylessNAS(input_tensor, cfg, version, flag),
+    'mnasnet':lambda input_tensor, cfg, version, flag: MnasNet(input_tensor, cfg, version, flag),
+    'squeezenet':lambda input_tensor, cfg, version, flag: SqueezeNet(input_tensor, cfg, version, flag),
+    'googlenet':lambda input_tensor, cfg, version, flag: GoogleNet(input_tensor, cfg, version, flag),
+    'densenet':lambda input_tensor, cfg, version, flag: DenseNet(input_tensor, cfg, version, flag),
+}
+
+
+class ModelGenerator:
+    def __init__(self, family_config, save_path):
+        random.seed(100)
+        self.save_path = save_path
+        self.log = {}
+        with open(family_config, "r") as f:
+            self.family_config = yaml.load(f, Loader = yaml.FullLoader)
+        self.log_file = os.path.join(self.save_path, self.family_config['model_family'] + '-log.json')
+
+    def add_to_log(self):
+        self.f = open(self.log_file, 'w')
+        self.f.write(json.dumps(self.log, indent=4))
+        self.f.flush()
+
+    def run(self):
+        c, h, w = self.family_config['input_shape']
+        sample_count = self.family_config['sample_count']
+
+        if self.family_config['model_family'] not in MODELZOO:
+            raise NotImplementedError
+
+        versions = [None]
+        if 'modelids' in self.family_config:
+            versions = self.family_config['modelids']
+
+        sconfigs = []
+        for vs in versions:
+            count_index = 0
+            while count_index < sample_count:
+                # generate the model with default configs from this family as the first model
+                sampling = False if count_index == 0 else True
+
+                input_tensor = keras.Input(shape=[h, w, c], batch_size=1)
+                model = MODELZOO[self.family_config['model_family']](input_tensor, self.family_config, vs, sampling)
+
+                if model.sconfig not in sconfigs:
+                    # save the h5 model file
+                    modelname = self.family_config['model_family']
+                    if vs: modelname += str(vs)
+                    output_tensor = model.out
+                    keras_model = keras.Model(input_tensor, output_tensor)
+                    save_name = modelname + "_" + str(count_index)
+                    keras_model.save(os.path.join(self.save_path, save_name + ".h5"))
+
+                    # log the model config
+                    sconfigs.append(model.sconfig)
+                    self.log[modelname + "_" + str(count_index)] = model.config
+                    self.add_to_log()
+                    count_index += 1
+
+# nasbench201
+def build_tiny_net(config_dict):
+    import subprocess
+    subprocess.check_output(f'python3 build_tiny_net.py -i "{config_dict["arch_str"]}" ' +
+                            f'-o {config_dict["pb_file_name"]} ' + (f'-t {config_dict["tflite_file_name"]}' if config_dict["tflite_file_name"] !=  '' else '')
+                            , shell = True, stderr = open(os.devnull, 'w'))

nn_meter/dataset/generator/interface.py

@@ -0,0 +1,34 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+def generate_dataset(model_faimly_list, num_of_samples=2000, output_folder=None, output_tflite_folder='', num_of_thread=12, input_nasbench201_descriptor=''):
+
+    # nasbench201_descriptor = json.loads(open(args.input_nasbench201_descriptor, 'r').read())
+    # nasbench201_acc_seq = sorted(nasbench201_descriptor, key = lambda x: nasbench201_descriptor[x]['acc'], reverse = True)
+
+    # nasbench201_acc_seq = nasbench201_acc_seq[:args.num_of_samples]
+
+    # #Pack args
+    # build_tiny_net_args = []
+    # for nasbench_keys in nasbench201_acc_seq:
+    #     arch_str = nasbench201_descriptor[nasbench_keys]['config']['arch_str']
+    #     pb_file_name = os.path.abspath(os.path.join(args.output_folder, 'nasbench201_%s.pb' % nasbench_keys))
+    #     if args.output_tflite_folder !=  '':
+    #         tflite_file_name = os.path.abspath(os.path.join(args.output_tflite_folder, 'nasbench201_%s.tflite' % nasbench_keys))
+    #     else:
+    #         tflite_file_name = ''
+    #     build_tiny_net_args.append({'arch_str': arch_str, 'pb_file_name': pb_file_name, 'tflite_file_name': tflite_file_name})
+
+    # if not os.path.isdir(args.output_folder):
+    #     os.mkdir(args.output_folder)
+
+    # if args.output_tflite_folder !=  '':
+    #     if not os.path.isdir(args.output_tflite_folder):
+    #         os.mkdir(args.output_tflite_folder)
+
+    # with multiprocessing.Pool(processes = args.num_of_thread) as p:
+    #     for _ in tqdm.tqdm(p.imap_unordered(build_tiny_net, build_tiny_net_args), total = len(build_tiny_net_args)):
+    #             pass
+    pass
+
+

nn_meter/dataset/generator/networks/tf_network/__init__.py

@@ -0,0 +1,14 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+from .alexnet import AlexNet
+from .mobilenetv1 import MobileNetV1
+from .mobilenetv2 import MobileNetV2
+from .mobilenetv3 import MobileNetV3
+from .shufflenetv2 import ShuffleNetV2
+from .resnetv1 import ResNetV1
+from .proxylessnas import ProxylessNAS
+from .mnasnet import MnasNet
+from .vgg import VGG
+from .squeezenet import SqueezeNet
+from .googlenet import GoogleNet
+from .densenet import DenseNet