moment-timeseries-foundation-model / moment

MOMENT: A Family of Open Time-series Foundation Models
https://moment-timeseries-foundation-model.github.io/
MIT License
319 stars 51 forks source link

ONNX Conversion error #43

Open kishoreb4 opened 1 month ago

kishoreb4 commented 1 month ago

I was able follow the script and finetune the model of a new dataset for forecasting.However, when i convert to onnx i'm getting the below error. Any help would be appreciated. Thanks in advance.

im getting the below error when in the code torch.onnx.export(model, (timeseries, input_mask), onnx_file_path)


RuntimeError Traceback (most recent call last) Cell In[59], line 2 1 # Convert the model to ONNX ----> 2 torch.onnx.export(model, (timeseries, input_mask), onnx_file_path)

File /libraries/gpu_py_3_11/lib/python3.11/site-packages/torch/onnx/utils.py:516, in export(model, args, f, export_params, verbose, training, input_names, output_names, operator_export_type, opset_version, do_constant_folding, dynamic_axes, keep_initializers_as_inputs, custom_opsets, export_modules_as_functions, autograd_inlining) 189 @_beartype.beartype 190 def export( 191 model: Union[torch.nn.Module, torch.jit.ScriptModule, torch.jit.ScriptFunction], (...) 208 autograd_inlining: Optional[bool] = True, 209 ) -> None: 210 r"""Exports a model into ONNX format. 211 212 If model is not a :class:torch.jit.ScriptModule nor a (...) 513 All errors are subclasses of :class:errors.OnnxExporterError. 514 """ --> 516 _export( 517 model, 518 args, 519 f, 520 export_params, 521 verbose, 522 training, 523 input_names, 524 output_names, 525 operator_export_type=operator_export_type, 526 opset_version=opset_version, 527 do_constant_folding=do_constant_folding, 528 dynamic_axes=dynamic_axes, 529 keep_initializers_as_inputs=keep_initializers_as_inputs, 530 custom_opsets=custom_opsets, 531 export_modules_as_functions=export_modules_as_functions, 532 autograd_inlining=autograd_inlining, 533 )

File /libraries/gpu_py_3_11/lib/python3.11/site-packages/torch/onnx/utils.py:1612, in _export(model, args, f, export_params, verbose, training, input_names, output_names, operator_export_type, export_type, opset_version, do_constant_folding, dynamic_axes, keep_initializers_as_inputs, fixed_batch_size, custom_opsets, add_node_names, onnx_shape_inference, export_modules_as_functions, autograd_inlining) 1609 dynamic_axes = {} 1610 _validate_dynamic_axes(dynamic_axes, model, input_names, output_names) -> 1612 graph, params_dict, torch_out = _model_to_graph( 1613 model, 1614 args, 1615 verbose, 1616 input_names, 1617 output_names, 1618 operator_export_type, 1619 val_do_constant_folding, 1620 fixed_batch_size=fixed_batch_size, 1621 training=training, 1622 dynamic_axes=dynamic_axes, 1623 ) 1625 # TODO: Don't allocate a in-memory string for the protobuf 1626 defer_weight_export = ( 1627 export_type is not _exporter_states.ExportTypes.PROTOBUF_FILE 1628 )

File /libraries/gpu_py_3_11/lib/python3.11/site-packages/torch/onnx/utils.py:1134, in _model_to_graph(model, args, verbose, input_names, output_names, operator_export_type, do_constant_folding, _disable_torch_constant_prop, fixed_batch_size, training, dynamic_axes) 1131 args = (args,) 1133 model = _pre_trace_quant_model(model, args) -> 1134 graph, params, torch_out, module = _create_jit_graph(model, args) 1135 params_dict = _get_named_param_dict(graph, params) 1137 try:

File /libraries/gpu_py_3_11/lib/python3.11/site-packages/torch/onnx/utils.py:1010, in _create_jit_graph(model, args) 1005 graph = _C._propagate_and_assign_input_shapes( 1006 graph, flattened_args, param_count_list, False, False 1007 ) 1008 return graph, params, torch_out, None -> 1010 graph, torch_out = _trace_and_get_graph_from_model(model, args) 1011 _C._jit_pass_onnx_lint(graph) 1012 state_dict = torch.jit._unique_state_dict(model)

File /libraries/gpu_py_3_11/lib/python3.11/site-packages/torch/onnx/utils.py:914, in _trace_and_get_graph_from_model(model, args) 912 prev_autocast_cache_enabled = torch.is_autocast_cache_enabled() 913 torch.set_autocast_cache_enabled(False) --> 914 trace_graph, torch_out, inputs_states = torch.jit._get_trace_graph( 915 model, 916 args, 917 strict=False, 918 _force_outplace=False, 919 _return_inputs_states=True, 920 ) 921 torch.set_autocast_cache_enabled(prev_autocast_cache_enabled) 923 warn_on_static_input_change(inputs_states)

File /libraries/gpu_py_3_11/lib/python3.11/site-packages/torch/jit/_trace.py:1310, in _get_trace_graph(f, args, kwargs, strict, _force_outplace, return_inputs, _return_inputs_states) 1308 if not isinstance(args, tuple): 1309 args = (args,) -> 1310 outs = ONNXTracedModule( 1311 f, strict, _force_outplace, return_inputs, _return_inputs_states 1312 )(*args, **kwargs) 1313 return outs

File /libraries/gpu_py_3_11/lib/python3.11/site-packages/torch/nn/modules/module.py:1532, in Module._wrapped_call_impl(self, *args, kwargs) 1530 return self._compiled_call_impl(*args, *kwargs) # type: ignore[misc] 1531 else: -> 1532 return self._call_impl(args, kwargs)

File /libraries/gpu_py_3_11/lib/python3.11/site-packages/torch/nn/modules/module.py:1541, in Module._call_impl(self, *args, *kwargs) 1536 # If we don't have any hooks, we want to skip the rest of the logic in 1537 # this function, and just call forward. 1538 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks 1539 or _global_backward_pre_hooks or _global_backward_hooks 1540 or _global_forward_hooks or _global_forward_pre_hooks): -> 1541 return forward_call(args, **kwargs) 1543 try: 1544 result = None

File /libraries/gpu_py_3_11/lib/python3.11/site-packages/torch/jit/_trace.py:138, in ONNXTracedModule.forward(self, *args) 135 else: 136 return tuple(out_vars) --> 138 graph, out = torch._C._create_graph_by_tracing( 139 wrapper, 140 in_vars + module_state, 141 _create_interpreter_name_lookup_fn(), 142 self.strict, 143 self._force_outplace, 144 ) 146 if self._return_inputs: 147 return graph, outs[0], ret_inputs[0]

File /libraries/gpu_py_3_11/lib/python3.11/site-packages/torch/jit/_trace.py:129, in ONNXTracedModule.forward..wrapper(args) 127 if self._return_inputs_states: 128 inputs_states.append(_unflatten(in_args, in_desc)) --> 129 outs.append(self.inner(trace_inputs)) 130 if self._return_inputs_states: 131 inputs_states[0] = (inputs_states[0], trace_inputs)

File /libraries/gpu_py_3_11/lib/python3.11/site-packages/momentfm/models/moment.py:222, in MOMENT.call(self, *args, kwargs) 221 def call(self, *args, *kwargs) -> TimeseriesOutputs: --> 222 return self.forward(args, kwargs)

File /libraries/gpu_py_3_11/lib/python3.11/site-packages/momentfm/models/moment.py:549, in MOMENT.forward(self, x_enc, mask, input_mask, kwargs) 547 return self.embed(x_enc=x_enc, input_mask=input_mask, kwargs) 548 elif self.task_name == TASKS.FORECASTING: --> 549 return self.forecast(x_enc=x_enc, input_mask=input_mask, kwargs) 550 elif self.task_name == TASKS.CLASSIFICATION: 551 return self.classify(x_enc=x_enc, input_mask=input_mask, kwargs)

File /libraries/gpu_py_3_11/lib/python3.11/site-packages/momentfm/models/moment.py:423, in MOMENT.forecast(self, x_enc, input_mask, *kwargs) 418 enc_in = enc_in.reshape( 419 (batch_size n_channels, n_patches, self.config.d_model) 420 ) 422 patch_view_mask = Masking.convert_seq_to_patch_view(input_mask, self.patch_len) --> 423 attention_mask = patch_view_mask.repeat_interleave(n_channels, dim=0) 424 outputs = self.encoder(inputs_embeds=enc_in, attention_mask=attention_mask) 425 enc_out = outputs.last_hidden_state

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument index in method wrapper_CUDA__index_select)