加载/保存状态字典
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,
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)
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()