Segmentation학습 > Onnx > C# - (3)export(Binary)

binary segmentation export : binary_segm_onnx_export.py

 

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
import torch
import cv2
import matplotlib.pyplot as plt
import albumentations as albu
import segmentation_models_pytorch as smp

ENCODER = 'se_resnext50_32x4d'
ENCODER_WEIGHTS = 'imagenet'
CLASSES = ['car']
ACTIVATION = 'sigmoid' # could be None for logits or 'softmax2d' for multiclass segmentation
DEVICE = 'cuda'

model = smp.FPN(
    encoder_name=ENCODER, 
    encoder_weights=ENCODER_WEIGHTS, 
    classes=len(CLASSES), 
    activation=ACTIVATION,
)

#model.load_state_dict(torch.load('./best_model_state_dict.pth'))
m = torch.load('./best_model_state_dict.pth')
model.load_state_dict(m)


batch_size = 1
#img = torch.randn(batch_size, 3, 896, 1280, requires_grad=True)
img = torch.randn(batch_size, 3, 384, 480, requires_grad=True)


torch.onnx.export(model, img, 
                 "./segmbin_best_mode.onnx", 
                  verbose=False,
                  opset_version=12,
                  input_names=['images'],
                  output_names=['output'],
                  do_constant_folding=True
                  #if device == 'cpu' else False
                  ,
                  dynamic_axes={'images': {0: 'batch_size'},
                                'output': {0: 'batch_size'}}
                 )