I am currently using a LSTM model to do some binary classification on a text dataset and was wondering how to go about extending this model to perform multi-label classification. The current model is as follows:
class LSTM(nn.Module):
def __init__(self, vocab_size, embedding_dim, hidden_dim, output_dim, n_layers, bidirectional,
dropout_rate, pad_index):
super().__init__()
self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=pad_index)
self.lstm = nn.LSTM(embedding_dim, hidden_dim, n_layers, bidirectional=bidirectional,
dropout=dropout_rate, batch_first=True)
self.fc = nn.Linear(hidden_dim * 2 if bidirectional else hidden_dim, output_dim)
self.dropout = nn.Dropout(dropout_rate)
def forward(self, ids, length):
embedded = self.dropout(self.embedding(ids))
packed_embedded = nn.utils.rnn.pack_padded_sequence(embedded, length, batch_first=True,
enforce_sorted=False)
packed_output, (hidden, cell) = self.lstm(packed_embedded)
output, output_length = nn.utils.rnn.pad_packed_sequence(packed_output)
if self.lstm.bidirectional:
hidden = self.dropout(torch.cat([hidden[-1], hidden[-2]], dim=-1))
else:
hidden = self.dropout(hidden[-1])
prediction = self.fc(hidden)
return prediction
Currently, the output of the model is 2 dimensions as the labels can either be 1 or 0. But I would be looking to convert it to classify a one hot encoding of the form [0, 0, 1, 0, 0, 1, 0]. I am aware I would probably need to change the criterion from nn.crossEntropyLoss but I am unsure what other steps to take.