def test_onnx_from_ckpt(self, tmpdir_fn, params):
        """Check that creating the onnx export from a model directly loaded from a checkpoint work as expected"""
        model_cls, clean = params
        model = model_cls(
            input_chunk_length=4, output_chunk_length=2, n_epochs=1, **tfm_kwargs_dev
        )
        onnx_filename = f"test_onnx_{model.model_name}.onnx"
        onnx_filename2 = f"test_onnx_{model.model_name}_weights.onnx"
        ckpt_filename = f"test_ckpt_{model.model_name}.pt"

        model.fit(
            series=self.ts_tg_with_static
            if model.supports_static_covariates
            else self.ts_tg,
            past_covariates=self.ts_pc if model.supports_past_covariates else None,
            future_covariates=self.ts_fc if model.supports_future_covariates else None,
        )
        model.save(ckpt_filename, clean=clean)

        # load the entire checkpoint
        load_kwargs = tfm_kwargs_dev if clean else {}
        model_loaded = model_cls.load(ckpt_filename, **load_kwargs)
        pred = model_loaded.predict(
            n=2,
            series=self.ts_tg_with_static
            if model_loaded.uses_static_covariates
            else self.ts_tg,
            past_covariates=self.ts_pc if model_loaded.uses_past_covariates else None,
            future_covariates=self.ts_fc
            if model_loaded.uses_future_covariates
            else None,
        )

        # export the loaded model
        model_loaded.to_onnx(onnx_filename)

        # onnx model loading and inference
        onnx_pred = self._helper_onnx_inference(
            model=model_loaded,
            onnx_filename=onnx_filename,
            series=self.ts_tg_with_static,
            past_covariates=self.ts_pc,
            future_covariates=self.ts_fc,
        )[0][0]

        # check that the predictions are similar
        assert pred.shape == onnx_pred.shape, "forecasts don't have the same shape."
        np.testing.assert_array_almost_equal(onnx_pred, pred.all_values(), decimal=4)

        # load only the weights
        model_weights = model_cls(
            input_chunk_length=4, output_chunk_length=2, n_epochs=1, **tfm_kwargs_dev
        )
        model_weights.load_weights(ckpt_filename)
        pred_weights = model_weights.predict(
            n=2,
            series=self.ts_tg_with_static
            if model_weights.uses_static_covariates
            else self.ts_tg,
            past_covariates=self.ts_pc if model_weights.uses_past_covariates else None,
            future_covariates=self.ts_fc
            if model_weights.uses_future_covariates
            else None,
        )

        # export the loaded model
        model_weights.to_onnx(onnx_filename2)

        # onnx model loading and inference
        onnx_pred_weights = self._helper_onnx_inference(
            model=model_weights,
            onnx_filename=onnx_filename2,
            series=self.ts_tg_with_static,
            past_covariates=self.ts_pc,
            future_covariates=self.ts_fc,
        )[0][0]

        assert pred_weights.shape == onnx_pred_weights.shape, (
            "forecasts don't have the same shape."
        )
        np.testing.assert_array_almost_equal(
            onnx_pred_weights, pred_weights.all_values(), decimal=4
        )
Check that creating the onnx export from a model directly loaded from a checkpoint work as expected
darts/tests/optional_deps/test_onnx.py