Dev gzf exp (#1654)

zhifu gao

2024-04-24 861147c7308b91068ffa02724fdf74ee623a909e

 funasr/models/ct_transformer/export_meta.py

@@ -37,31 +37,31 @@
    y = self.decoder(h)
    return y


def export_dummy_inputs(self):
    length = 120
    text_indexes = torch.randint(0, self.embed.num_embeddings, (2, length)).type(torch.int32)
    text_lengths = torch.tensor([length-20, length], dtype=torch.int32)
    return (text_indexes, text_lengths)


def export_input_names(self):
    return ['inputs', 'text_lengths']
    return ["inputs", "text_lengths"]


def export_output_names(self):
    return ['logits']
    return ["logits"]


def export_dynamic_axes(self):
    return {
        'inputs': {
            0: 'batch_size',
            1: 'feats_length'
        "inputs": {0: "batch_size", 1: "feats_length"},
        "text_lengths": {
            0: "batch_size",
        },
        'text_lengths': {
            0: 'batch_size',
        },
        'logits': {
            0: 'batch_size',
            1: 'logits_length'
        },
        "logits": {0: "batch_size", 1: "logits_length"},
    }


def export_name(self):
    return "model.onnx"