-
Notifications
You must be signed in to change notification settings - Fork 25.7k
Description
OS: WIN10
CUDA 10
CUDNN 7.3
PYTORCH 1.0
I trained unet. It works fine in a python environment.
I want to make it work in c++ by libtorch.
first,I use Official example to test.
Like this:
import torch
import torchvision
model = torchvision.models.resnet50(pretrained=True)
model.eval()
example = torch.zeros(1, 3, 224, 224)
traced_script_module = torch.jit.trace(model, example)
traced_script_module.save("model.pt")
then run it in cpp,like this:
torch::DeviceType device_type = at::kCUDA;
std::cout << "use cuda" << std::endl;
std::shared_ptr<torch::jit::script::Module> module = torch::jit::load(model_file);
module->to(device_type);
assert(module != nullptr);
std::cout << "ok\n";
std::vector<torch::jit::IValue> inputs;
inputs.push_back(torch::ones({1, 3, 224, 224}).to(at::kCUDA));
at::Tensor output = module->forward(inputs).toTensor();
std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/5) << '\n';
get right reslut , all thing look like so perfect.
but when I use my unet to Instead of resnet, the error occur.it will accident exit without any useful message.only show "Not loaded kerenlbase.pdb"
I just use this code to convert my UNET model.
import torch
import torch.nn.functional as F
from unet import UNet
model = UNet(n_channels=3, n_classes=1)
model_path = "./MODEL.pth"
model.load_state_dict(torch.load(model_path))
model.eval()
example = torch.zeros(1, 3, 224, 224)
traced_script_module = torch.jit.trace(model, example)
traced_script_module.save("model.pt")
convert unet model is a little bit different of resnet. one load from model zoo,other load by my local network and model.
cpp is exactly the same.
Who can explain it? thank you.