| import numpy as np | |
| import ncnn | |
| import torch | |
| def test_inference(): | |
| torch.manual_seed(0) | |
| in0 = torch.rand(1, 3, 384, 384, dtype=torch.float) | |
| in1 = torch.rand(1, 3, 384, 384, dtype=torch.float) | |
| out = [] | |
| with ncnn.Net() as net: | |
| net.load_param("/Users/raoulritter/STB-VMM/20x/modelpnnx20x.ncnn.param") | |
| net.load_model("/Users/raoulritter/STB-VMM/20x/modelpnnx20x.ncnn.bin") | |
| with net.create_extractor() as ex: | |
| ex.input("in0", ncnn.Mat(in0.squeeze(0).numpy()).clone()) | |
| ex.input("in1", ncnn.Mat(in1.squeeze(0).numpy()).clone()) | |
| _, out0 = ex.extract("out0") | |
| out.append(torch.from_numpy(np.array(out0)).unsqueeze(0)) | |
| _, out1 = ex.extract("out1") | |
| out.append(torch.from_numpy(np.array(out1)).unsqueeze(0)) | |
| _, out2 = ex.extract("out2") | |
| out.append(torch.from_numpy(np.array(out2)).unsqueeze(0)) | |
| _, out3 = ex.extract("out3") | |
| out.append(torch.from_numpy(np.array(out3))) | |
| if len(out) == 1: | |
| return out[0] | |
| else: | |
| return tuple(out) | |