Plugin to launch PyTorch models

Development Pipeline & Plugins Plugins
1

I made this plugin to launch script modules trained in PyTorch. PyTorch has C++ frontend which can work without Python, and this is great, but, as I figured out, it conflicts with UE4. So I added a wrapper dll loaded in runtime, and now you can launch your trained neural nets from UE4.

Link: GitHub - AntiAnti/SimplePyTorch: UE4 Plugin to execute trained PyTorch modules

Note: the only supported platform is Win64.

1 Like
2

wow, thanks a lot for this!!

3

Hey @YuriNK @Diomedes246 is there a sample implementation of this? Im interested in bringing in MiDaS pytorch model for media depth estimation (after which i can estimate the normal map so video textures can be relit)

4

Download release of the plugin from github, it already contains all pytorch DLLs inside.

I have implementation in my project which I can’t share, but here is a short sample. Note: you need to be very accurate with size and shape of tensors. Any error cause crash.

//////////////////////////////
// header
//////////////////////////////

protected:

	UPROPERTY()
	FSimpleTorchTensor FacialEmotionsInData;
	
	UPROPERTY()
	FSimpleTorchTensor FacialEmotionsOutData;
	
	UPROPERTY()
	USimpleTorchModule* FacialEmotionsModel;
	
	UPROPERTY()
	bool bTorchModelReady;

//////////////////////////////
// cpp
//////////////////////////////

// BeginPlay
{
	bTorchModelReady = false;

	// initialize PyTorch jit script model
	FacialEmotionsModel = USimpleTorchModule::CreateSimpleTorchModule(this);
	if (FacialEmotionsModel)
	{
		FString FileName = FPaths::ProjectDir() / TEXT("emotions.tmod");
		if (FPaths::FileExists(FileName) && FacialEmotionsModel->LoadTorchScriptModel(FileName))
		{
			FacialEmotionsInData.Create({ 1 });
			FacialEmotionsOutData.Create({ 1, NeuralNetOutputColumns.Num() );
			bTorchModelReady = true;
		}
		else
		{
			UE_LOG(LogVH, Warning, TEXT("Unable to load torch jit model from file (%s)"), *FileName);
		}
	}
	else
	{
		UE_LOG(LogVH, Warning, TEXT("Unable to create torch jit model wrapper"));
	}
}

// usage
{
	//...
	FacialEmotionsInData[0] = (float)SymbolNum;
	FacialEmotionsModel->ExecuteModelMethod(TEXT("eval_compute"), FacialEmotionsInData, FacialEmotionsOutData);
	check(FacialEmotionsOutData.GetDataSize() == Columns.Num());
	
	for (int32 n = 0; n < Columns.Num(); n++)
	{
		float val = FMath::Clamp(FacialEmotionsOutData[0][n].Item(), 0.f, 1.f);		
		UE_LOG(LogVH, Log, TEXT("Result value [%d] = %f"), n+1, val);
	}

	//...	
}
5

In Python:

class EmotionsNet(nn.Module):
    # ...

    @torch.jit.export
    def eval_compute(self, x):
        #...

def nnsave(Net):
    Net.eval()
    scripted_module = torch.jit.script(Net)
    torch.jit.save(scripted_module, __script_path + '/emotions.tmod')