Make ONNXmodel

[big-chan]

import torch
from model import SSD300
from apex.parallel.LARC import LARC
from apex import amp
from apex.parallel import DistributedDataParallel as DDP
from apex.fp16_utils import *
def load_model_weight(model, model_path):
    model_checkpoint=torch.load(model_path)["model"]
    model.load_state_dict(model_checkpoint.state_dict())
    #model = amp.initialize(model, opt_level='O2')
    model = model.eval()
    #model=model.to("cpu")
    return model
def export_onnx_model(model, input_shape, onnx_path, input_names=None, output_names=None, dynamic_axes=None):
    inputs_R = torch.ones((1,3,512,640)).to("cuda")
    inputs_T = torch.ones((1,1,512,640)).to("cuda")
    #import pdb;pdb.set_trace()
    #model(inputs)
    torch.onnx.export(model, (inputs_R,inputs_T), onnx_path, input_names=input_names, output_names=output_names)
    
if __name__ == "__main__":

    checkpoint ='./checkpoint_ssd300.pth.tar083'

    model_path = "model_checkpoint.pth"
    model = SSD300(2)
    model=model.to("cuda")
    model = load_model_weight(model, checkpoint)

    input_shape = (1, 1, 512, 640)
    onnx_path = "halfway.onnx"
    # input_names=['input']
    # output_names=['output']
    # dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}}
    export_onnx_model(model, input_shape, onnx_path)

it was based this site

[big-chan]

주의 사항

• pooling 시 ceil_mode를 false 로 해주세요
• pow 를 없애주세요