pytorch KAIR example on Android

Viewed 134

I stuck trying to trace/scipt ffdnet KAIR's model to android. Model's forward looks like:

def forward(self, x): #, paddingBottom, paddingRight): #, sigma):

    noise_level_model = 15
    sigma = torch.full((1, 1, 1, 1), noise_level_model / 255.).type_as(x)

    h, w = x.size()[-2:]
    paddingBottom = int(np.ceil(h/2)*2-h)
    paddingRight = int(np.ceil(w/2)*2-w)

    x = torch.nn.ReplicationPad2d((0, paddingRight, 0, paddingBottom))(x)

    x = self.m_down(x)
    # m = torch.ones(sigma.size()[0], sigma.size()[1], x.size()[-2], x.size()[-1]).type_as(x).mul(sigma)
    m = sigma.repeat(1, 1, x.size()[-2], x.size()[-1])
    x = torch.cat((x, m), 1)
    x = self.model(x)
    x = self.m_up(x)
    
    x = x[..., :h, :w]
    return x

If I trace that I get some warnings about padding arguments but model works on Android. Problem is that it isn't work with input of different sizes, only size same as 'test1.jpeg':

model_name = 'ffdnet_color'
model_pool = 'model_zoo'
model_path = os.path.join(model_pool, model_name + '.pth')
n_channels = 3
nc = 96
nb = 12

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

model = net(in_nc=n_channels, out_nc=n_channels, nc=nc, nb=nb, act_mode='R')
model.load_state_dict(torch.load(model_path), strict=True)
model.eval()
for k, v in model.named_parameters():
    v.requires_grad = False
model = model.to(device)

img = 'testsets/myset/test1.jpeg'
img_name, ext = os.path.splitext(os.path.basename(img))
img_L = util.imread_uint(img, n_channels=n_channels)
img_L = util.uint2single(img_L)
noise_level_model = 15
img_L = util.single2tensor4(img_L)
img_L = img_L.to(device)
sigma_ = torch.full((1, 1, 1, 1), noise_level_model / 255)
sigma = torch.full((1, 1, 1, 1), noise_level_model / 255.).type_as(img_L)

traced_model = torch.jit.trace(model, img_L)
traced_optimized = optimize_for_mobile(traced_model)

save_path = os.path.splitext(os.path.basename(model_path))[0] + '-mobile.pth'
traced_optimized.save(save_path)

I've tried to script model with traced_model = torch.jit.script(model) but got errors:

TypeError: cannot create weak reference to 'numpy.ufunc' object

What should I do to achieve model works with different input size on mobile?

1 Answers

I encountered a similar issue. It was due to my model using numpy math operations (which are numpy.ufunc). I fixed the issue by replacing all of numpy ufuncs (i.e. np.add, np.ceil, and +, - etc on ndarrays) with corresponding torch versions (i.e. torch.add, torch.sub etc).

Related