nagadomi / waifu2x

Image Super-Resolution for Anime-Style Art
http://waifu2x.udp.jp/
MIT License
27.49k stars 2.71k forks source link

transform the model to onnx #382

Open xsy123xsy opened 3 years ago

xsy123xsy commented 3 years ago

I am trying to test the model in a Real Time video process engine, before that I need to transform the model into onnx Can I use a python script like this:

import re
import os
import torch
from utils.network_srmd import SRMD as net

n_channels = 3            # fixed
nc = 128                  # fixed, number of channels
nb = 12                   # fixed, number of conv layers

model_list = ['srmd_x2', 'srmd_x3', 'srmd_x4', 'srmdnf_x2', 'srmdnf_x3', 'srmdnf_x4']
model_pool = 'model_zoo'
output_model_pool = 'onnx_models'

for model_name in model_list:
    sf = [int(s) for s in re.findall(r'\d+', model_name)][0]  # scale factor
    in_nc = 18 if 'nf' in model_name else 19
    model_path = os.path.join(model_pool, model_name+'.pth')

    model = net(in_nc=in_nc, out_nc=n_channels, nc=nc, nb=nb, upscale=sf, act_mode='R', upsample_mode='pixelshuffle')
    model.load_state_dict(torch.load(model_path), strict=False)
    model.eval()
    for k, v in model.named_parameters():
        v.requires_grad = False
    x = torch.randn((1, 3, 1, 1))
    k_pca = torch.randn(1, 15, 1, 1)
    m = k_pca.repeat(1, 1, x.size()[-2], x.size()[-1])
    x = torch.cat((x, m), 1)
    if 'nf' not in model_name:  # noise-free SR
        noise_level = torch.zeros(1, 1, 1, 1)
        x = torch.cat((x, noise_level), 1)
    torch_out = torch.onnx._export(model, x, os.path.join(output_model_pool, model_name+".onnx"), export_params=True)
    print("Convert "+ model_name+" success!")
nagadomi commented 3 years ago

I haven't tried ONNX, but the model converted to pytorch and the example of image conversion(tiled rendering) with pytorch can be found in https://github.com/nagadomi/nunif (You need to install git-lfs to pull the pretrained model files)

model parameter/state conversion script: https://github.com/nagadomi/nunif/blob/master/nunif/cli/import_waifu2x_models.py

nagadomi commented 3 years ago

Also, waifu2x-caffe has a model converted to caffe model.

xsy123xsy commented 3 years ago

Thanks for help, I will try to transform the model from pytorch to onnx. By the way, to do the transformation, can I use python script like this: https://github.com/HollowMan6/SRMD-Pytorch/blob/main/pytorch2onnx.py

------------------ 原始邮件 ------------------ 发件人: "nagadomi/waifu2x" <notifications@github.com>; 发送时间: 2021年2月6日(星期六) 上午6:02 收件人: "nagadomi/waifu2x"<waifu2x@noreply.github.com>; 抄送: "xsy_xsy"<366366235@qq.com>;"Author"<author@noreply.github.com>; 主题: Re: [nagadomi/waifu2x] transform the model to onnx (#382)

I haven't tried ONNX, but the model converted to pytorch and the example of image conversion(tiled rendering) with pytorch can be found in https://github.com/nagadomi/nunif (You need to install git-lfs to pull the pretrained model files)

model parameter/state conversion script: https://github.com/nagadomi/nunif/blob/master/nunif/cli/import_waifu2x_models.py

— You are receiving this because you authored the thread. Reply to this email directly, view it on GitHub, or unsubscribe.