Skip to content
Snippets Groups Projects
Commit 79e13e56 authored by chenjun2hao's avatar chenjun2hao
Browse files

onnx with cpu

parent 1daf52a3
No related branches found
No related tags found
No related merge requests found
......@@ -27,7 +27,7 @@ def parse_args():
parser.add_argument('--cfg',
help='experiment configure file name',
default="experiments/cityscapes/ddrnet23_slim.yaml",
default="experiments/cityscapes/ddrnet23.yaml",
type=str)
parser.add_argument('opts',
help="Modify config options using the command-line",
......@@ -81,7 +81,7 @@ def main():
# model_state_file = os.path.join(final_output_dir, 'final_state.pth')
logger.info('=> loading model from {}'.format(model_state_file))
pretrained_dict = torch.load(model_state_file)
pretrained_dict = torch.load(model_state_file, map_location='cpu')
if 'state_dict' in pretrained_dict:
pretrained_dict = pretrained_dict['state_dict']
model_dict = model.state_dict()
......@@ -100,8 +100,7 @@ def main():
x = torch.randn((1,3,480,640))
torch_out = net(x)
# output_path = "output/tensorrt/resnet50/resnet50_bilinear.onnx"
output_path = "output/ddrnet23_slim.onnx"
output_path = "output/ddrnet23.onnx"
torch.onnx.export(net, # model being run
x, # model input (or a tuple for multiple inputs)
output_path, # where to save the model (can be a file or file-like object)
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment