{ "cells": [ { "cell_type": "markdown", "id": "1c296cb7", "metadata": {}, "source": [ "# TMVA_SOFIE_PyTorch\n", "This macro provides a simple example for the parsing of PyTorch .pt file\n", "into RModel object and further generating the .hxx header files for inference.\n", "\n", "\n", "\n", "**Author:** Sanjiban Sengupta \n", "This notebook tutorial was automatically generated with ROOTBOOK-izer from the macro found in the ROOT repository on Tuesday, May 19, 2026 at 08:23 PM." ] }, { "cell_type": "code", "execution_count": 1, "id": "2a2ad553", "metadata": { "collapsed": false, "execution": { "iopub.execute_input": "2026-05-19T20:23:35.105250Z", "iopub.status.busy": "2026-05-19T20:23:35.105111Z", "iopub.status.idle": "2026-05-19T20:23:35.701202Z", "shell.execute_reply": "2026-05-19T20:23:35.685714Z" } }, "outputs": [], "source": [ "using namespace TMVA::Experimental;\n", "\n", "TString pythonSrc = \"\\\n", "import torch\\n\\\n", "import torch.nn as nn\\n\\\n", "\\n\\\n", "model = nn.Sequential(\\n\\\n", " nn.Linear(32,16),\\n\\\n", " nn.ReLU(),\\n\\\n", " nn.Linear(16,8),\\n\\\n", " nn.ReLU()\\n\\\n", " )\\n\\\n", "\\n\\\n", "criterion = nn.MSELoss()\\n\\\n", "optimizer = torch.optim.SGD(model.parameters(),lr=0.01)\\n\\\n", "\\n\\\n", "x=torch.randn(2,32)\\n\\\n", "y=torch.randn(2,8)\\n\\\n", "\\n\\\n", "for i in range(500):\\n\\\n", " y_pred = model(x)\\n\\\n", " loss = criterion(y_pred,y)\\n\\\n", " optimizer.zero_grad()\\n\\\n", " loss.backward()\\n\\\n", " optimizer.step()\\n\\\n", "\\n\\\n", "model.eval()\\n\\\n", "m = torch.jit.script(model)\\n\\\n", "torch.jit.save(m,'PyTorchModel.pt')\\n\";" ] }, { "cell_type": "markdown", "id": "ca43ac30", "metadata": {}, "source": [ "Running the Python script to generate PyTorch .pt file" ] }, { "cell_type": "code", "execution_count": 2, "id": "4bccc772", "metadata": { "collapsed": false, "execution": { "iopub.execute_input": "2026-05-19T20:23:35.735126Z", "iopub.status.busy": "2026-05-19T20:23:35.734965Z", "iopub.status.idle": "2026-05-19T20:23:49.775132Z", "shell.execute_reply": "2026-05-19T20:23:49.767520Z" } }, "outputs": [], "source": [ "TMacro m;\n", "m.AddLine(pythonSrc);\n", "m.SaveSource(\"make_pytorch_model.py\");\n", "gSystem->Exec(\"python3 make_pytorch_model.py\");" ] }, { "cell_type": "markdown", "id": "59951007", "metadata": {}, "source": [ "Parsing a PyTorch model requires the shape and data-type of input tensor\n", "Data-type of input tensor defaults to Float if not specified" ] }, { "cell_type": "code", "execution_count": 3, "id": "5444abf2", "metadata": { "collapsed": false, "execution": { "iopub.execute_input": "2026-05-19T20:23:49.784366Z", "iopub.status.busy": "2026-05-19T20:23:49.784229Z", "iopub.status.idle": "2026-05-19T20:23:50.040655Z", "shell.execute_reply": "2026-05-19T20:23:50.019243Z" } }, "outputs": [], "source": [ "std::vector inputTensorShapeSequential{2, 32};\n", "std::vector> inputShapesSequential{inputTensorShapeSequential};" ] }, { "cell_type": "markdown", "id": "982a8d08", "metadata": {}, "source": [ "Parsing the saved PyTorch .pt file into RModel object" ] }, { "cell_type": "code", "execution_count": 4, "id": "14b94398", "metadata": { "collapsed": false, "execution": { "iopub.execute_input": "2026-05-19T20:23:50.059285Z", "iopub.status.busy": "2026-05-19T20:23:50.059135Z", "iopub.status.idle": "2026-05-19T20:23:52.063742Z", "shell.execute_reply": "2026-05-19T20:23:52.041525Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Torch Version: 2.12.0+cpu" ] }, { "name": "stdout", "output_type": "stream", "text": [ "\n" ] } ], "source": [ "SOFIE::RModel model = SOFIE::PyTorch::Parse(\"PyTorchModel.pt\", inputShapesSequential);" ] }, { "cell_type": "markdown", "id": "f545e846", "metadata": {}, "source": [ "Generating inference code" ] }, { "cell_type": "code", "execution_count": 5, "id": "e2e5beed", "metadata": { "collapsed": false, "execution": { "iopub.execute_input": "2026-05-19T20:23:52.069041Z", "iopub.status.busy": "2026-05-19T20:23:52.068791Z", "iopub.status.idle": "2026-05-19T20:23:52.333062Z", "shell.execute_reply": "2026-05-19T20:23:52.318648Z" } }, "outputs": [], "source": [ "model.Generate();\n", "model.OutputGenerated(\"PyTorchModel.hxx\");" ] }, { "cell_type": "markdown", "id": "06ecd36d", "metadata": {}, "source": [ "Printing required input tensors" ] }, { "cell_type": "code", "execution_count": 6, "id": "f226d96b", "metadata": { "collapsed": false, "execution": { "iopub.execute_input": "2026-05-19T20:23:52.335323Z", "iopub.status.busy": "2026-05-19T20:23:52.335086Z", "iopub.status.idle": "2026-05-19T20:23:52.583817Z", "shell.execute_reply": "2026-05-19T20:23:52.569327Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "\n", "\n", "Model requires following inputs:\n", "Fully Specified Tensor name: input1\ttype: float\tshape: [2,32]\n", "\n" ] } ], "source": [ "std::cout << \"\\n\\n\";\n", "model.PrintRequiredInputTensors();" ] }, { "cell_type": "markdown", "id": "e91d91ba", "metadata": {}, "source": [ "Printing initialized tensors (weights)" ] }, { "cell_type": "code", "execution_count": 7, "id": "845107ed", "metadata": { "collapsed": false, "execution": { "iopub.execute_input": "2026-05-19T20:23:52.589012Z", "iopub.status.busy": "2026-05-19T20:23:52.588870Z", "iopub.status.idle": "2026-05-19T20:23:52.814721Z", "shell.execute_reply": "2026-05-19T20:23:52.806340Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "\n", "\n", "Model initialized the following tensors:\n", "Tensor name: \"2bias\"\ttype: float\tshape: [8]\n", "Tensor name: \"0weight\"\ttype: float\tshape: [16,32]\n", "Tensor name: \"2weight\"\ttype: float\tshape: [8,16]\n", "Tensor name: \"0bias\"\ttype: float\tshape: [16]\n", "\n" ] } ], "source": [ "std::cout << \"\\n\\n\";\n", "model.PrintInitializedTensors();" ] }, { "cell_type": "markdown", "id": "1e937010", "metadata": {}, "source": [ "Printing intermediate tensors" ] }, { "cell_type": "code", "execution_count": 8, "id": "b78f2f19", "metadata": { "collapsed": false, "execution": { "iopub.execute_input": "2026-05-19T20:23:52.816504Z", "iopub.status.busy": "2026-05-19T20:23:52.816352Z", "iopub.status.idle": "2026-05-19T20:23:53.049452Z", "shell.execute_reply": "2026-05-19T20:23:53.048771Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "\n", "\n", "Model specify the following intermediate tensors:\n", "Tensor name: \"result3\"\ttype: float\tshape: [2,8]\n", "Tensor name: \"result\"\ttype: float\tshape: [2,16]\n", "Tensor name: \"input2\"\ttype: float\tshape: [2,8]\n", "Tensor name: \"input0\"\ttype: float\tshape: [2,16]\n", "\n" ] } ], "source": [ "std::cout << \"\\n\\n\";\n", "model.PrintIntermediateTensors();" ] }, { "cell_type": "markdown", "id": "9491c3cc", "metadata": {}, "source": [ "Checking if tensor already exist in model" ] }, { "cell_type": "code", "execution_count": 9, "id": "5216918f", "metadata": { "collapsed": false, "execution": { "iopub.execute_input": "2026-05-19T20:23:53.051298Z", "iopub.status.busy": "2026-05-19T20:23:53.051178Z", "iopub.status.idle": "2026-05-19T20:23:53.264185Z", "shell.execute_reply": "2026-05-19T20:23:53.263528Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "\n", "\n", "Tensor \"0weight\" already exist: true\n", "\n", "Shape of tensor \"0weight\": 16,32,\n", "\n", "Data type of tensor \"0weight\": float" ] } ], "source": [ "std::cout << \"\\n\\nTensor \\\"0weight\\\" already exist: \" << std::boolalpha << model.CheckIfTensorAlreadyExist(\"0weight\")\n", " << \"\\n\\n\";\n", "std::vector tensorShape = model.GetTensorShape(\"0weight\");\n", "std::cout << \"Shape of tensor \\\"0weight\\\": \";\n", "for (auto &it : tensorShape) {\n", " std::cout << it << \",\";\n", "}\n", " std::cout<<\"\\n\\nData type of tensor \\\"0weight\\\": \";\n", " SOFIE::ETensorType tensorType = model.GetTensorType(\"0weight\");\n", " std::cout<\n", "#include \n", "#include \"TMVA/SOFIE_common.hxx\"\n", "#include \n", "\n", "namespace TMVA_SOFIE_PyTorchModel{\n", "namespace BLAS{\n", "\textern \"C\" void sgemv_(const char * trans, const int * m, const int * n, const float * alpha, const float * A,\n", "\t const int * lda, const float * X, const int * incx, const float * beta, const float * Y, const int * incy);\n", "\textern \"C\" void sgemm_(const char * transa, const char * transb, const int * m, const int * n, const int * k,\n", "\t const float * alpha, const float * A, const int * lda, const float * B, const int * ldb,\n", "\t const float * beta, float * C, const int * ldc);\n", "}//BLAS\n", "struct Session;\n", "inline void doInfer(Session const &session, float const* tensor_input1, float *tensor_result3 );\n", "struct Session {\n", "// initialized (weights and constant) tensors\n", "std::vector fTensor_2bias = std::vector(8);\n", "float * tensor_2bias = fTensor_2bias.data();\n", "std::vector fTensor_0weight = std::vector(512);\n", "float * tensor_0weight = fTensor_0weight.data();\n", "std::vector fTensor_2weight = std::vector(128);\n", "float * tensor_2weight = fTensor_2weight.data();\n", "std::vector fTensor_0bias = std::vector(16);\n", "float * tensor_0bias = fTensor_0bias.data();\n", "\n", "//--- Allocating session memory pool to be used for allocating intermediate tensors\n", "std::vector fIntermediateMemoryPool = std::vector(256);\n", "\n", "\n", "// --- Positioning intermediate tensor memory --\n", " // Allocating memory for intermediate tensor input0 with size 128 bytes\n", "float* tensor_input0 = reinterpret_cast(fIntermediateMemoryPool.data() + 0);\n", "\n", " // Allocating memory for intermediate tensor result with size 128 bytes\n", "float* tensor_result = reinterpret_cast(fIntermediateMemoryPool.data() + 128);\n", "\n", " // Allocating memory for intermediate tensor input2 with size 64 bytes\n", "float* tensor_input2 = reinterpret_cast(fIntermediateMemoryPool.data() + 64);\n", "\n", " // Allocating memory for intermediate tensor result3 with size 64 bytes\n", "float* tensor_result3 = reinterpret_cast(fIntermediateMemoryPool.data() + 0);\n", "\n", "\n", "Session(std::string filename =\"PyTorchModel.dat\") {\n", "\n", "//--- reading weights from file\n", " std::ifstream f;\n", " f.open(filename);\n", " if (!f.is_open()) {\n", " throw std::runtime_error(\"tmva-sofie failed to open file \" + filename + \" for input weights\");\n", " }\n", " using TMVA::Experimental::SOFIE::ReadTensorFromStream;\n", " ReadTensorFromStream(f, tensor_2bias, \"tensor_2bias\", 8);\n", " ReadTensorFromStream(f, tensor_0weight, \"tensor_0weight\", 512);\n", " ReadTensorFromStream(f, tensor_2weight, \"tensor_2weight\", 128);\n", " ReadTensorFromStream(f, tensor_0bias, \"tensor_0bias\", 16);\n", " f.close();\n", "\n", "}\n", "\n", "\n", "\n", "std::vector infer(float const* tensor_input1){\n", " std::vector output_tensor_result3(16);\n", " doInfer(*this, tensor_input1, output_tensor_result3.data() );\n", " return {output_tensor_result3};\n", "}\n", "}; // end of Session\n", "\n", "\n", "// Input tensor dimensions\n", "using TMVA::Experimental::SOFIE::SingleDim;\n", "using TMVA::Experimental::SOFIE::TensorDims;\n", "using TMVA::Experimental::SOFIE::makeDims;\n", "\n", "constexpr std::array dim_input1{SingleDim{2}, SingleDim{32}};\n", "\n", "constexpr std::array inputTensorDims{\n", " makeDims(dim_input1)\n", "};\n", "\n", "constexpr bool hasDynamicInputTensors{false};\n", "\n", "\n", "// Output tensor dimensions\n", "constexpr std::array dim_result3{SingleDim{2}, SingleDim{8}};\n", "\n", "constexpr std::array outputTensorDims{\n", " makeDims(dim_result3)\n", "};\n", "\n", "constexpr bool hasDynamicOutputTensors{false};\n", "\n", "inline void doInfer(Session const &session, float const* tensor_input1, float *tensor_result3 ) {\n", "\n", " auto &tensor_0bias = session.tensor_0bias;\n", " auto &tensor_0weight = session.tensor_0weight;\n", " auto &tensor_2bias = session.tensor_2bias;\n", " auto &tensor_2weight = session.tensor_2weight;\n", " auto &tensor_input0 = session.tensor_input0;\n", " auto &tensor_input2 = session.tensor_input2;\n", " auto &tensor_result = session.tensor_result;\n", "\n", "\n", "//--------- Gemm op_0 { 2 , 32 } * { 16 , 32 } -> { 2 , 16 }\n", " for (size_t j = 0; j < 2; j++) { \n", " size_t y_index = 16 * j;\n", " TMVA::Experimental::SOFIE::Copy(tensor_input0 + y_index, tensor_0bias, 16);\n", " }\n", " TMVA::Experimental::SOFIE::Gemm_Call(tensor_input0, true, false, 16, 2, 32, 1, tensor_0weight, tensor_input1, 1,nullptr);\n", "\n", "//------ RELU\n", " for (int id = 0; id < 32 ; id++){\n", " tensor_result[id] = ((tensor_input0[id] > 0 )? tensor_input0[id] : 0);\n", " }\n", "\n", "//--------- Gemm op_2 { 2 , 16 } * { 8 , 16 } -> { 2 , 8 }\n", " for (size_t j = 0; j < 2; j++) { \n", " size_t y_index = 8 * j;\n", " TMVA::Experimental::SOFIE::Copy(tensor_input2 + y_index, tensor_2bias, 8);\n", " }\n", " TMVA::Experimental::SOFIE::Gemm_Call(tensor_input2, true, false, 8, 2, 16, 1, tensor_2weight, tensor_result, 1,nullptr);\n", "\n", "//------ RELU\n", " for (int id = 0; id < 16 ; id++){\n", " tensor_result3[id] = ((tensor_input2[id] > 0 )? tensor_input2[id] : 0);\n", " }\n", "\n", "}\n", "} //TMVA_SOFIE_PyTorchModel\n", "\n", "#endif // ROOT_TMVA_SOFIE_PYTORCHMODEL\n" ] } ], "source": [ " std::cout<<\"\\n\\n\";\n", " model.PrintGenerated();" ] } ], "metadata": { "kernelspec": { "display_name": "ROOT C++", "language": "c++", "name": "root" }, "language_info": { "codemirror_mode": "text/x-c++src", "file_extension": ".C", "mimetype": " text/x-c++src", "name": "c++" } }, "nbformat": 4, "nbformat_minor": 5 }