使用DEV PyTr火炬1将PyTrac模型加载到C++中

如果您可以使用Pytorch hub的UNET模型。它将与TorchScript一起工作。

import torch

# downloading the model from torchhub
model = torch.hub.load('mateuszbuda/brain-segmentation-pytorch', 'unet',
    in_channels=3, out_channels=1, init_features=32, pretrained=True)

#  downloading the sample
import urllib
url, filename = ("https://github.com/mateuszbuda/brain-segmentation-pytorch/raw/master/assets/TCGA_CS_4944.png", "TCGA_CS_4944.png")
try: urllib.URLopener().retrieve(url, filename)
except: urllib.request.urlretrieve(url, filename)
    
# reading the sample and some prerequisites for transformation
import numpy as np
from PIL import Image
from torchvision import transforms

input_image = Image.open(filename)

m, s = np.mean(input_image, axis=(0, 1)), np.std(input_image, axis=(0, 1))
preprocess = transforms.Compose([transforms.ToTensor(),transforms.Normalize(mean=m, std=s),])

input_tensor = preprocess(input_image)

input_batch = input_tensor.unsqueeze(0)

# creating the trace
traced_module = torch.jit.trace(model,input_batch)

# running the trace
traced_module(input_batch)

PS:torch.jit.trace/torch.jit.script都不支持所有的torch功能,所以在外部库中使用它们总是很棘手的。