Spaces:
Runtime error
Runtime error
#!/usr/bin/env python3 | |
# -*- coding:utf-8 -*- | |
import argparse | |
import time | |
import sys | |
import os | |
import torch | |
import torch.nn as nn | |
import onnx | |
ROOT = os.getcwd() | |
if str(ROOT) not in sys.path: | |
sys.path.append(str(ROOT)) | |
from yolov6.models.yolo import * | |
from yolov6.models.effidehead import Detect | |
from yolov6.layers.common import * | |
from yolov6.utils.events import LOGGER | |
from yolov6.utils.checkpoint import load_checkpoint | |
if __name__ == '__main__': | |
parser = argparse.ArgumentParser() | |
parser.add_argument('--weights', type=str, default='./yolov6s.pt', help='weights path') | |
parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size') # height, width | |
parser.add_argument('--batch-size', type=int, default=1, help='batch size') | |
parser.add_argument('--half', action='store_true', help='FP16 half-precision export') | |
parser.add_argument('--inplace', action='store_true', help='set Detect() inplace=True') | |
parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0, 1, 2, 3 or cpu') | |
args = parser.parse_args() | |
args.img_size *= 2 if len(args.img_size) == 1 else 1 # expand | |
print(args) | |
t = time.time() | |
# Check device | |
cuda = args.device != 'cpu' and torch.cuda.is_available() | |
device = torch.device('cuda:0' if cuda else 'cpu') | |
assert not (device.type == 'cpu' and args.half), '--half only compatible with GPU export, i.e. use --device 0' | |
# Load PyTorch model | |
model = load_checkpoint(args.weights, map_location=device, inplace=True, fuse=True) # load FP32 model | |
for layer in model.modules(): | |
if isinstance(layer, RepVGGBlock): | |
layer.switch_to_deploy() | |
# Input | |
img = torch.zeros(args.batch_size, 3, *args.img_size).to(device) # image size(1,3,320,192) iDetection | |
# Update model | |
if args.half: | |
img, model = img.half(), model.half() # to FP16 | |
model.eval() | |
for k, m in model.named_modules(): | |
if isinstance(m, Conv): # assign export-friendly activations | |
if isinstance(m.act, nn.SiLU): | |
m.act = SiLU() | |
elif isinstance(m, Detect): | |
m.inplace = args.inplace | |
y = model(img) # dry run | |
# ONNX export | |
try: | |
LOGGER.info('\nStarting to export ONNX...') | |
export_file = args.weights.replace('.pt', '.onnx') # filename | |
torch.onnx.export(model, img, export_file, verbose=False, opset_version=12, | |
training=torch.onnx.TrainingMode.EVAL, | |
do_constant_folding=True, | |
input_names=['image_arrays'], | |
output_names=['outputs'], | |
) | |
# Checks | |
onnx_model = onnx.load(export_file) # load onnx model | |
onnx.checker.check_model(onnx_model) # check onnx model | |
LOGGER.info(f'ONNX export success, saved as {export_file}') | |
except Exception as e: | |
LOGGER.info(f'ONNX export failure: {e}') | |
# Finish | |
LOGGER.info('\nExport complete (%.2fs)' % (time.time() - t)) | |