This macro provides a simple example for the parsing of PyTorch .pt file into RModel object and further generating the .hxx header files for inference.
Author: Sanjiban Sengupta
This notebook tutorial was automatically generated with ROOTBOOK-izer from the macro found in the ROOT repository on Monday, March 27, 2023 at 09:59 AM.
using namespace TMVA::Experimental;
TString pythonSrc = "\
import torch\n\
import torch.nn as nn\n\
\n\
model = nn.Sequential(\n\
nn.Linear(32,16),\n\
nn.ReLU(),\n\
nn.Linear(16,8),\n\
nn.ReLU()\n\
)\n\
\n\
criterion = nn.MSELoss()\n\
optimizer = torch.optim.SGD(model.parameters(),lr=0.01)\n\
\n\
x=torch.randn(2,32)\n\
y=torch.randn(2,8)\n\
\n\
for i in range(500):\n\
y_pred = model(x)\n\
loss = criterion(y_pred,y)\n\
optimizer.zero_grad()\n\
loss.backward()\n\
optimizer.step()\n\
\n\
model.eval()\n\
m = torch.jit.script(model)\n\
torch.jit.save(m,'PyTorchModel.pt')\n";
Running the Python script to generate PyTorch .pt file
TMVA::PyMethodBase::PyInitialize();
TMacro m;
m.AddLine(pythonSrc);
m.SaveSource("make_pytorch_model.py");
gSystem->Exec(TMVA::Python_Executable() + " make_pytorch_model.py");
Parsing a PyTorch model requires the shape and data-type of input tensor Data-type of input tensor defaults to Float if not specified
std::vector<size_t> inputTensorShapeSequential{2,32};
std::vector<std::vector<size_t>> inputShapesSequential{inputTensorShapeSequential};
Parsing the saved PyTorch .pt file into RModel object
SOFIE::RModel model = SOFIE::PyTorch::Parse("PyTorchModel.pt",inputShapesSequential);
Torch Version: 1.13.1+cu117
Generating inference code
model.Generate();
model.OutputGenerated("PyTorchModel.hxx");
Printing required input tensors
std::cout<<"\n\n";
model.PrintRequiredInputTensors();
Model requires following inputs: Fully Specified Tensor name: input1 type: float shape: [2,32]
Printing initialized tensors (weights)
std::cout<<"\n\n";
model.PrintInitializedTensors();
Model initialized the following tensors: Tensor name: "2bias" type: float shape: [8] Tensor name: "0weight" type: float shape: [16,32] Tensor name: "2weight" type: float shape: [8,16] Tensor name: "0bias" type: float shape: [16]
Printing intermediate tensors
std::cout<<"\n\n";
model.PrintIntermediateTensors();
Model specify the following intermediate tensors: Tensor name: "result3" type: float shape: [2,8] Tensor name: "2biasbcast" type: float shape: [2,8] Tensor name: "input2" type: float shape: [2,8] Tensor name: "input0" type: float shape: [2,16] Tensor name: "result" type: float shape: [2,16] Tensor name: "0biasbcast" type: float shape: [2,16]
Checking if tensor already exist in model
std::cout<<"\n\nTensor \"0weight\" already exist: "<<std::boolalpha<<model.CheckIfTensorAlreadyExist("0weight")<<"\n\n";
std::vector<size_t> tensorShape = model.GetTensorShape("0weight");
std::cout<<"Shape of tensor \"0weight\": ";
for(auto& it:tensorShape){
std::cout<<it<<",";
}
std::cout<<"\n\nData type of tensor \"0weight\": ";
SOFIE::ETensorType tensorType = model.GetTensorType("0weight");
std::cout<<SOFIE::ConvertTypeToString(tensorType);
Tensor "0weight" already exist: true Shape of tensor "0weight": 16,32, Data type of tensor "0weight": float
Printing generated inference code
std::cout<<"\n\n";
model.PrintGenerated();
//Code generated automatically by TMVA for Inference of Model file [PyTorchModel.pt] at [Mon Mar 27 09:59:13 2023]
#ifndef TMVA_SOFIE_PYTORCHMODEL
#define TMVA_SOFIE_PYTORCHMODEL
#include<algorithm>
#include<vector>
#include "TMVA/SOFIE_common.hxx"
#include <fstream>
namespace TMVA_SOFIE_PyTorchModel{
namespace BLAS{
extern "C" void sgemv_(const char * trans, const int * m, const int * n, const float * alpha, const float * A,
const int * lda, const float * X, const int * incx, const float * beta, const float * Y, const int * incy);
extern "C" void sgemm_(const char * transa, const char * transb, const int * m, const int * n, const int * k,
const float * alpha, const float * A, const int * lda, const float * B, const int * ldb,
const float * beta, float * C, const int * ldc);
}//BLAS
struct Session {
std::vector<float> fTensor_2bias = std::vector<float>(8);
float * tensor_2bias = fTensor_2bias.data();
std::vector<float> fTensor_0weight = std::vector<float>(512);
float * tensor_0weight = fTensor_0weight.data();
std::vector<float> fTensor_2weight = std::vector<float>(128);
float * tensor_2weight = fTensor_2weight.data();
std::vector<float> fTensor_0bias = std::vector<float>(16);
float * tensor_0bias = fTensor_0bias.data();
std::vector<float> fTensor_result3 = std::vector<float>(16);
float * tensor_result3 = fTensor_result3.data();
std::vector<float> fTensor_2biasbcast = std::vector<float>(16);
float * tensor_2biasbcast = fTensor_2biasbcast.data();
std::vector<float> fTensor_input2 = std::vector<float>(16);
float * tensor_input2 = fTensor_input2.data();
std::vector<float> fTensor_input0 = std::vector<float>(32);
float * tensor_input0 = fTensor_input0.data();
std::vector<float> fTensor_result = std::vector<float>(32);
float * tensor_result = fTensor_result.data();
std::vector<float> fTensor_0biasbcast = std::vector<float>(32);
float * tensor_0biasbcast = fTensor_0biasbcast.data();
Session(std::string filename ="") {
if (filename.empty()) filename = "PyTorchModel.dat";
std::ifstream f;
f.open(filename);
if (!f.is_open()){
throw std::runtime_error("tmva-sofie failed to open file for input weights");
}
std::string tensor_name;
int length;
f >> tensor_name >> length;
if (tensor_name != "tensor_2bias" ) {
std::string err_msg = "TMVA-SOFIE failed to read the correct tensor name; expected name is tensor_2bias , read " + tensor_name;
throw std::runtime_error(err_msg);
}
if (length != 8) {
std::string err_msg = "TMVA-SOFIE failed to read the correct tensor size; expected size is 8 , read " + std::to_string(length) ;
throw std::runtime_error(err_msg);
}
for (int i =0; i < length; ++i)
f >> tensor_2bias[i];
f >> tensor_name >> length;
if (tensor_name != "tensor_0weight" ) {
std::string err_msg = "TMVA-SOFIE failed to read the correct tensor name; expected name is tensor_0weight , read " + tensor_name;
throw std::runtime_error(err_msg);
}
if (length != 512) {
std::string err_msg = "TMVA-SOFIE failed to read the correct tensor size; expected size is 512 , read " + std::to_string(length) ;
throw std::runtime_error(err_msg);
}
for (int i =0; i < length; ++i)
f >> tensor_0weight[i];
f >> tensor_name >> length;
if (tensor_name != "tensor_2weight" ) {
std::string err_msg = "TMVA-SOFIE failed to read the correct tensor name; expected name is tensor_2weight , read " + tensor_name;
throw std::runtime_error(err_msg);
}
if (length != 128) {
std::string err_msg = "TMVA-SOFIE failed to read the correct tensor size; expected size is 128 , read " + std::to_string(length) ;
throw std::runtime_error(err_msg);
}
for (int i =0; i < length; ++i)
f >> tensor_2weight[i];
f >> tensor_name >> length;
if (tensor_name != "tensor_0bias" ) {
std::string err_msg = "TMVA-SOFIE failed to read the correct tensor name; expected name is tensor_0bias , read " + tensor_name;
throw std::runtime_error(err_msg);
}
if (length != 16) {
std::string err_msg = "TMVA-SOFIE failed to read the correct tensor size; expected size is 16 , read " + std::to_string(length) ;
throw std::runtime_error(err_msg);
}
for (int i =0; i < length; ++i)
f >> tensor_0bias[i];
f.close();
{
float * data = TMVA::Experimental::SOFIE::UTILITY::UnidirectionalBroadcast<float>(tensor_0bias,{ 16 }, { 2 , 16 });
std::copy(data, data + 32, tensor_0biasbcast);
delete [] data;
}
{
float * data = TMVA::Experimental::SOFIE::UTILITY::UnidirectionalBroadcast<float>(tensor_2bias,{ 8 }, { 2 , 8 });
std::copy(data, data + 16, tensor_2biasbcast);
delete [] data;
}
}
std::vector<float> infer(float* tensor_input1){
//--------- Gemm
char op_0_transA = 'n';
char op_0_transB = 't';
int op_0_m = 2;
int op_0_n = 16;
int op_0_k = 32;
float op_0_alpha = 1;
float op_0_beta = 1;
int op_0_lda = 32;
int op_0_ldb = 32;
std::copy(tensor_0biasbcast, tensor_0biasbcast + 32, tensor_input0);
BLAS::sgemm_(&op_0_transB, &op_0_transA, &op_0_n, &op_0_m, &op_0_k, &op_0_alpha, tensor_0weight, &op_0_ldb, tensor_input1, &op_0_lda, &op_0_beta, tensor_input0, &op_0_n);
//------ RELU
for (int id = 0; id < 32 ; id++){
tensor_result[id] = ((tensor_input0[id] > 0 )? tensor_input0[id] : 0);
}
//--------- Gemm
char op_2_transA = 'n';
char op_2_transB = 't';
int op_2_m = 2;
int op_2_n = 8;
int op_2_k = 16;
float op_2_alpha = 1;
float op_2_beta = 1;
int op_2_lda = 16;
int op_2_ldb = 16;
std::copy(tensor_2biasbcast, tensor_2biasbcast + 16, tensor_input2);
BLAS::sgemm_(&op_2_transB, &op_2_transA, &op_2_n, &op_2_m, &op_2_k, &op_2_alpha, tensor_2weight, &op_2_ldb, tensor_result, &op_2_lda, &op_2_beta, tensor_input2, &op_2_n);
//------ RELU
for (int id = 0; id < 16 ; id++){
tensor_result3[id] = ((tensor_input2[id] > 0 )? tensor_input2[id] : 0);
}
std::vector<float> ret (tensor_result3, tensor_result3 + 16);
return ret;
}
};
} //TMVA_SOFIE_PyTorchModel
#endif // TMVA_SOFIE_PYTORCHMODEL