pytorchのモデルをJITでコンパイル

JIT_compile.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable

from models import modelName

import numpy as np
import cv2

batchsize =1
input_channel = 3
output_channel= 3
input_height = input_width = 256



model = modelName(input_channel, output_channel).cuda()
model.load_state_dict(torch.load('./model.pth')) #  pytorchモデルを読む
model.eval()

example = torch.rand(1, 3, 256, 256)  #入力用に定義

device = torch.device("cuda")
mode = "cuda"
model = torch.jit.trace(model, example.to(device))
model.save("Net_h{}_w{}_{}.pt".format(input_height, input_width, mode))

print("Net_h{}_w{}_{}.pt is exported".format(output_channel, input_height, input_width, mode))

input_x_np = np.zeros((batchsize, input_channel, input_height, input_width)).astype(np.float32)