简体   繁体   English

如何加载预训练的 pytorch 权重

[英]How to load pretrained pytorch weight for

I am following this blog https://pytorch.org/tutorials/advanced/super_resolution_with_onnxruntime.html want to run pytorch model in onnx runtime. I am following this blog https://pytorch.org/tutorials/advanced/super_resolution_with_onnxruntime.html want to run pytorch model in onnx runtime. here in example it has given a pretrained weight a URL how to load a pretrained weight from local disk.在此示例中,它给出了预训练权重 URL 如何从本地磁盘加载预训练权重。

# Load pretrained model weights
model_url = 'https://s3.amazonaws.com/pytorch/test_data/export/superres_epoch100-44c6958e.pth'
path = "/content/best.pt"
batch_size = 1    # just a random number

# Initialize model with the pretrained weights
map_location = lambda storage, loc: storage
if torch.cuda.is_available():
    map_location = None
torch_model.load_state_dict(model_zoo.load_url(model_url, map_location=map_location))

# set the model to inference mode
torch_model.eval()

I want to load the weight which is defined as Path.我想加载定义为路径的权重。

If you want to load the state dict from a path, this is what you should do:如果您想从路径加载 state 字典,您应该这样做:

torch_model.load_state_dict(torch.load(path))

This should work.这应该有效。

声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM