| | |
| | | is_onnx = kwargs.get("type", "onnx") == "onnx" |
| | | encoder_class = tables.encoder_classes.get(kwargs["encoder"] + "Export") |
| | | model.encoder = encoder_class(model.encoder, onnx=is_onnx) |
| | | |
| | | |
| | | model.forward = types.MethodType(export_forward, model) |
| | | model.export_dummy_inputs = types.MethodType(export_dummy_inputs, model) |
| | | model.export_input_names = types.MethodType(export_input_names, model) |
| | |
| | | |
| | | return model |
| | | |
| | | |
| | | def export_forward(self, feats: torch.Tensor, *args, **kwargs): |
| | | |
| | | |
| | | scores, out_caches = self.encoder(feats, *args) |
| | | |
| | | |
| | | return scores, out_caches |
| | | |
| | | |
| | | def export_dummy_inputs(self, data_in=None, frame=30): |
| | | if data_in is None: |
| | | speech = torch.randn(1, frame, self.encoder_conf.get("input_dim")) |
| | | else: |
| | | speech = None # Undo |
| | | |
| | | speech = None # Undo |
| | | |
| | | cache_frames = self.encoder_conf.get("lorder") + self.encoder_conf.get("rorder") - 1 |
| | | in_cache0 = torch.randn(1, self.encoder_conf.get("proj_dim"), cache_frames, 1) |
| | | in_cache1 = torch.randn(1, self.encoder_conf.get("proj_dim"), cache_frames, 1) |
| | | in_cache2 = torch.randn(1, self.encoder_conf.get("proj_dim"), cache_frames, 1) |
| | | in_cache3 = torch.randn(1, self.encoder_conf.get("proj_dim"), cache_frames, 1) |
| | | |
| | | |
| | | return (speech, in_cache0, in_cache1, in_cache2, in_cache3) |
| | | |
| | | |
| | | def export_input_names(self): |
| | | return ['speech', 'in_cache0', 'in_cache1', 'in_cache2', 'in_cache3'] |
| | | return ["speech", "in_cache0", "in_cache1", "in_cache2", "in_cache3"] |
| | | |
| | | |
| | | def export_output_names(self): |
| | | return ['logits', 'out_cache0', 'out_cache1', 'out_cache2', 'out_cache3'] |
| | | return ["logits", "out_cache0", "out_cache1", "out_cache2", "out_cache3"] |
| | | |
| | | |
| | | def export_dynamic_axes(self): |
| | | return { |
| | | 'speech': { |
| | | 1: 'feats_length' |
| | | }, |
| | | "speech": {1: "feats_length"}, |
| | | } |
| | | |
| | | def export_name(self, ): |
| | | |
| | | def export_name( |
| | | self, |
| | | ): |
| | | return "model.onnx" |