I try to run libtorch in iOS, but get a weird output. I use the below python code to export the model to file. When load in python it works well, But when load in ios the output looks weird.
export code:
class TestModule(nn.Module):
def __init__(self):
super(TestModule, self).__init__()
self.op = nn.Sequential(
nn.Conv2d(1, 2, 3, padding=1),
)
self._init_weights()
def _init_weights(self):
for m in self.modules():
if isinstance(m, nn.ConvTranspose2d) or isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight.data)
m.bias.data.zero_()
def forward(self, x):
return self.op(x)
down = TestModule()
script_model = torch.jit.trace(down, torch.rand(1, 1, 4, 4))
print(script_model)
script_model.save('model/down_4.pt')
Load in iOS:
auto module_ = torch::jit::load(modulePath.UTF8String);
at::Tensor input = at::ones({1, 1, 4, 4}, at::kFloat);
torch::autograd::AutoGradMode guard(false);
auto output = module_.forward({input}).toTensor();
std::cout << output << "\n";
iOS output:
libtorch version: 1.4.0