Feature: Support ONNX conversion for models

[ogencoglu]

Do you support ONNX conversion?

[LennartPurucker]

Currently, the code base does not have code to support ONNX conversion.

We are looking into supporting this and have tried some ideas. However, due to the new architecture, we still need to finalize this.

@AlexanderPfefferle do you have more information about this?

[AlexanderPfefferle]

#165 enabled ONNX export
The tests contain an example for how to export:

TabPFN/tests/test_classifier_interface.py

Lines 227 to 289 in b0bc461

	class ModelWrapper(nn.Module):
	def __init__(self, original_model): # noqa: D107
	super().__init__()
	self.model = original_model
	def forward(
	self,
	X,
	y,
	single_eval_pos,
	only_return_standard_out,
	categorical_inds,
	):
	return self.model(
	None,
	X,
	y,
	single_eval_pos=single_eval_pos,
	only_return_standard_out=only_return_standard_out,
	categorical_inds=categorical_inds,
	)
	@pytest.mark.filterwarnings("ignore::torch.jit.TracerWarning")
	def test_onnx_exportable_cpu(X_y: tuple[np.ndarray, np.ndarray]) -> None:
	if os.name == "nt":
	pytest.skip("onnx export is not tested on windows")
	X, y = X_y
	with torch.no_grad():
	classifier = TabPFNClassifier(n_estimators=1, device="cpu", random_state=42)
	# load the model so we can access it via classifier.model_
	classifier.fit(X, y)
	# this is necessary if cuda is available
	classifier.predict(X)
	# replicate the above call with random tensors of same shape
	X = torch.randn(
	(X.shape[0] * 2, 1, X.shape[1] + 1),
	generator=torch.Generator().manual_seed(42),
	)
	y = (
	torch.rand(y.shape, generator=torch.Generator().manual_seed(42))
	.round()
	.to(torch.float32)
	)
	dynamic_axes = {
	"X": {0: "num_datapoints", 1: "batch_size", 2: "num_features"},
	"y": {0: "num_labels"},
	}
	torch.onnx.export(
	ModelWrapper(classifier.model_).eval(),
	(X, y, y.shape[0], True, []),
	io.BytesIO(),
	input_names=[
	"X",
	"y",
	"single_eval_pos",
	"only_return_standard_out",
	"categorical_inds",
	],
	output_names=["output"],
	opset_version=17, # using 17 since we use torch>=2.1
	dynamic_axes=dynamic_axes,
	)

Currently, you can only use the exported model with the exact same number of features that the example input for exporting had, but @LeoGrin is looking into that.