Upload export.py

export.py

@@ -0,0 +1,156 @@
+import argparse
+from itertools import chain
+
+import torch
+import torch.nn as nn
+from transformers import LlamaConfig, DynamicCache
+
+from midi_model import MIDIModel, config_name_list, MIDIModelConfig
+
+
+class MIDIModelBase(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.net = model.net
+
+    def forward(self, x, past_kv):
+        cache = DynamicCache.from_legacy_cache(past_kv)
+        x = self.net.embed_tokens(x)
+        x = x.sum(dim=-2)
+        x = self.net.forward(inputs_embeds=x,
+                             past_key_values=cache,
+                             use_cache=True)
+        return x.last_hidden_state, cache.to_legacy_cache()
+
+
+class MIDIModelToken(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.net_token = model.net_token
+        self.lm_head = model.lm_head
+
+    def forward(self, hidden_state, x, past_kv):
+        cache = DynamicCache.from_legacy_cache(past_kv)
+        x = self.net_token.embed_tokens(x)
+        x = torch.cat([hidden_state, x], dim=1)
+        hidden_state = x
+        hidden_state = self.net_token.forward(inputs_embeds=hidden_state,
+                                              past_key_values=cache,
+                                              use_cache=True).last_hidden_state
+        return self.lm_head(hidden_state), cache.to_legacy_cache()
+
+
+def export_onnx(model, model_inputs, input_names, output_names, dynamic_axes, meta_data, path):
+    import onnx
+    from onnxsim import simplify
+    torch.onnx.export(model,  # model being run
+                      model_inputs,  # model input (or a tuple for multiple inputs)
+                      path,  # where to save the model (can be a file or file-like object)
+                      export_params=True,  # store the trained parameter weights inside the model file
+                      opset_version=14,  # the ONNX version to export the model to
+                      do_constant_folding=True,  # whether to execute constant folding for optimization
+                      input_names=input_names,  # the model's input names
+                      output_names=output_names,  # the model's output names
+                      verbose=True,
+                      dynamic_axes=dynamic_axes
+                      )
+    onnx_model = onnx.load(path)
+    model_simp, check = simplify(onnx_model)
+    assert check, "Simplified ONNX model could not be validated"
+    for k, v in meta_data.items():
+        meta = model_simp.metadata_props.add()
+        meta.key, meta.value = k, str(v)
+    onnx.save(model_simp, path)
+    print('finished exporting onnx')
+
+
+def get_past_kv(config: LlamaConfig, batch_size=1, past_seq_len=16, torch_dtype= torch.float32, device="cpu"):
+    head_size = config.hidden_size // config.num_attention_heads
+    past_kv = [
+        (
+            torch.rand(batch_size, config.num_attention_heads,
+                       past_seq_len, head_size, dtype=torch_dtype, device=device),
+            torch.rand(batch_size, config.num_attention_heads,
+                       past_seq_len, head_size, dtype=torch_dtype, device=device),
+        )
+        for _ in range(config.num_hidden_layers)
+    ]
+    input_names = list(
+        chain.from_iterable(
+            (f"past_key_values.{i}.key", f"past_key_values.{i}.value") for i in
+            range(config.num_hidden_layers)
+        )
+    )
+    output_names = list(
+        chain.from_iterable((f"present.{i}.key", f"present.{i}.value") for i in range(config.num_hidden_layers))
+    )
+    return past_kv, input_names, output_names
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--ckpt", type=str, default="model.ckpt", help="load ckpt"
+    )
+    parser.add_argument(
+        "--config", type=str, default="tv2o-medium", choices=config_name_list, help="model config"
+    )
+    parser.add_argument(
+        "--lora", type=str, default="", help="load lora"
+    )
+    parser.add_argument(
+        "--model-base-out", type=str, default="model_base.onnx", help="model base output path"
+    )
+    parser.add_argument(
+        "--model-token-out", type=str, default="model_token.onnx", help="model token output path"
+    )
+    opt = parser.parse_args()
+    config = MIDIModelConfig.from_name(opt.config)
+    tokenizer = config.tokenizer
+    model = MIDIModel(config).to(device="cpu")
+    ckpt = torch.load(opt.ckpt, map_location="cpu")
+    state_dict = ckpt.get("state_dict", ckpt)
+    model.load_state_dict(state_dict, strict=False)
+    if opt.lora != "":
+        model.load_merge_lora(opt.lora)
+    model.eval()
+    model_base = MIDIModelBase(model).eval()
+    model_token = MIDIModelToken(model).eval()
+    meta_data = {"config_name": opt.config, "config": config}
+    past_kv_shape = {0: "batch", 2: "past_seq"}
+    present_kv_shape = {0: "batch", 2: "present_seq"}
+    with torch.no_grad():
+        dynamic_axes = {
+            "x": {0: "batch", 1: "mid_seq", 2: "token_seq"},
+            "hidden": {0: "batch", 1: "mid_seq"}
+        }
+        x = torch.randint(tokenizer.vocab_size, (1, 16, tokenizer.max_token_seq), dtype=torch.int64, device="cpu")
+        past_kv, input_names, output_names= get_past_kv(config.net_config, past_seq_len=16,
+                                                        torch_dtype=torch.float32)
+        for name in input_names:
+            dynamic_axes[name] = past_kv_shape
+        for name in output_names:
+            dynamic_axes[name] = present_kv_shape
+        input_names = [ "x"] + input_names
+        output_names = ["hidden"] + output_names
+        export_onnx(model_base, (x, past_kv),
+                    input_names, output_names, dynamic_axes, meta_data, opt.model_base_out)
+
+        dynamic_axes = {
+            "x": {0: "batch", 1: "token_seq"},
+            "hidden": {0: "batch", 1: "states"},
+            "y": {0: "batch", 1: "token_seq1"}
+        }
+        hidden = torch.randn(1, 1, config.n_embd, device="cpu")
+        x = torch.randint(tokenizer.vocab_size, (1, tokenizer.max_token_seq //2), dtype=torch.int64, device="cpu")
+        past_kv, input_names, output_names = get_past_kv(config.net_token_config,
+                                                         past_seq_len=(tokenizer.max_token_seq // 2),
+                                                         torch_dtype=torch.float32)
+        for name in input_names:
+            dynamic_axes[name] = past_kv_shape
+        for name in output_names:
+            dynamic_axes[name] = present_kv_shape
+        input_names = ["hidden", "x"] + input_names
+        output_names = ["y"] + output_names
+        export_onnx(model_token, (hidden, x, past_kv),
+                    input_names, output_names, dynamic_axes, meta_data, opt.model_token_out)