def export2caffe(weights, num_classes, img_size):
model = UNet(num_classes)
weights = torch.load(weights, map_location='cpu')
model.load_state_dict(weights['model'])
model.eval()
fuse(model)
name = 'DeepLabV3Plus'
dummy_input = torch.ones([1, 3, img_size[1], img_size[0]])
pytorch2caffe.trans_net(model, dummy_input, name)
pytorch2caffe.save_prototxt('{}.prototxt'.format(name))
pytorch2caffe.save_caffemodel('{}.caffemodel'.format(name))
def export2caffe(weights, num_classes, img_size):
os.environ['MODEL_EXPORT'] = '1'
model = YOLOV3(num_classes)
weights = torch.load(weights, map_location='cpu')
model.load_state_dict(weights['model'])
model.eval()
fuse(model)
name = 'RYOLOV3'
dummy_input = torch.ones([1, 3, img_size[1], img_size[0]])
pytorch2caffe.trans_net(model, dummy_input, name)
pytorch2caffe.save_prototxt('{}.prototxt'.format(name))
pytorch2caffe.save_caffemodel('{}.caffemodel'.format(name))
def export2caffe(weights, num_classes, img_size):
model = MobileNetV2(num_classes)
weights = torch.load(weights, map_location='cpu')
model.load_state_dict(weights['model'])
model.eval()
fuse(model)
dummy_input = torch.ones([1, 3, img_size[1], img_size[0]])
torch.onnx.export(model,
dummy_input,
'MobileNetV2.onnx',
input_names=['input'],
output_names=['output'],
opset_version=7)