pytorch模型保存与加载

加载/保存状态字典

class RNN(nn.Module):

    def __init__(self, vocab_size, embedding_dim, hidden_dim):

        super(RNN, self).__init__()

        
        self.embedding = nn.Embedding(vocab_size, embedding_dim)        #单词数，嵌入向量维度
        self.embedding.weight.requires_grad = False
      
        self.rnn = nn.LSTM(embedding_dim, hidden_dim, num_layers=2,     #双向RNN，所以下面使用hidden_dim*2
                           bidirectional=True, dropout=0.5)
    
        self.fc = nn.Linear(hidden_dim*2, 2)
        self.dropout = nn.Dropout(0.5)


    def forward(self, x):
      
        embedding = self.dropout(self.embedding(x))

      
        input = embedding.transpose(0, 1)
        output, (hidden, cell) = self.rnn(input)          #[h0,c0]随机初始化

        
        hidden = torch.cat([hidden[-2], hidden[-1]], dim=1)   #双向，所以要把最后两个输出连接

        
        hidden = self.dropout(hidden)
        out = self.fc(hidden)

        return out

torch.save(rnn.state_dict(),'tx.pt')
rnn = RNN(10004, 50, 256)   
rnn.load_state_dict(torch.load('tx.pt'))
rnn.eval()   #测试模式

加载/保存整个模型

torch.save(rnn, 'txd.pt')
rnn1 = torch.load('txd.pt')
rnn1.eval()   #测试模式