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))
import torch
from torchvision.models import resnet
from pytorch2caffe import pytorch2caffe
name='resnet18'
resnet18=resnet.resnet18()
resnet18.eval()
dummy_input=torch.ones([1,3,224,224])
pytorch2caffe.trans_net(resnet18, dummy_input, name)
pytorch2caffe.save_prototxt('{}.prototxt'.format(name))
pytorch2caffe.save_caffemodel('{}.caffemodel'.format(name))