在 PyTorch 中使用预训练的权重

Question

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

这篇论文是学习在黑暗中看东西。

模型架构如下：

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()

然后是从谷歌驱动器/云存储导入训练的权重，并定义将训练的权重放入网络的函数。

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

Answer 1

如果您使用的是 google colab

#mount drive onto google colab

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

定义权重的路径

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

解压 tar 文件

!tar -xvf weights.pth.tar

将权重加载到模型网络中

net=torch.load(weights_path)