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)