在PyTorch中使用预先训练的权重

时间:2020-07-20 11:16:10

标签: deep-learning neural-network computer-vision pytorch

我正在PyTorch中实施基于计算机视觉的研究论文。我已经通过参考本文建立了模型架构。作者已将已保存的权重以“ .pth.tar”格式上传到GitHub。我想在模型中加入相同的权重,以便跳过训练和优化部分,直接从神经网络获取输出。

这篇论文正在学习如何在黑暗中看。

模型架构如下:

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv1d(32, 12, 1)
        .
        .  
   def forward(self, x):
        x = F.relu(self.conv1(x))
        .
        .    
        return x
net = Net()

然后是从Google驱动器/云存储导入训练后的权重并定义将训练后的权重放入网络的功能。

PS:两者的模型架构完全相同

1 个答案:

答案 0 :(得分:1)

如果您使用的是Google colab

<td><?php echo form_checkbox( 'leave_id[]',$Dleave->leave_id); ?></td>

定义权重的路径

#mount drive onto google colab

from google.colab import drive
drive.mount('/content/gdrive')

解压缩tar文件

weights_path="/content/gdrive/My Drive/weights.pth"

将权重加载到模型网络中

!tar -xvf weights.pth.tar