Stefano Piovesan
Stefano Piovesan

Reputation: 1215

Inference pytorch C++ with alexnet and cv::imread image

I am trying to infer with a C++ application an image classification task using an alexnet pre-trained net.
I have successfully inferred a dog image loading the net with python:

alexnet = torchvision.models.alexnet(pretrained=True)
img = Image.open("dog.jpg")
transform = transforms.Compose([
 transforms.Resize(256),                
 transforms.CenterCrop(224),        
 transforms.ToTensor(),                  
 transforms.Normalize(                   
 mean=[0.485, 0.456, 0.406],         
 std=[0.229, 0.224, 0.225]              
 )])
img_t = transform(img)
batch_t = torch.unsqueeze(img_t, 0)
alexnet.forward(batch_t)
_, index = torch.max(out, 1)

Result index is 208, Labrador_retriever, that looks good.
Then I save the net to be loaded from a C++ application

example = torch.rand(1, 3, 224, 224)
traced_script_module_alex = torch.jit.trace(alexnet, example)
traced_script_module.save("alexnet.pt")

When I load to C++, I get the wrong result:

cv::Mat img = cv::imread("dog.jpg");
cv::resize(img, img, cv::Size(224, 224), cv::INTER_CUBIC);

// Convert the image and label to a tensor.
torch::Tensor img_tensor = torch::from_blob(img.data, { 1, img.rows, img.cols, 3 }, torch::kByte);
img_tensor = img_tensor.permute({ 0, 3, 1, 2 }); // convert to CxHxW
img_tensor = img_tensor.to(torch::kFloat);
std::vector<torch::jit::IValue> input;
input.push_back(img_tensor);
torch::jit::script::Module  module = torch::jit::load("alexnet.pt");
at::Tensor output = module.forward(input).toTensor();
std::cout << output.argmax(1) << '\n';

the argmax is 463, bucket. I think I am not looking at the same image; what am I missing...?

Upvotes: 0

Views: 2001

Answers (1)

lenik
lenik

Reputation: 23556

Your C++ code is missing this part of your Python code:

transform = transforms.Compose([
 transforms.Resize(256),                
 transforms.CenterCrop(224),        
 transforms.ToTensor(),                  
 transforms.Normalize(                   
 mean=[0.485, 0.456, 0.406],         
 std=[0.229, 0.224, 0.225]              
 )])
img_t = transform(img)

Upvotes: 1

Related Questions