110 lines
4.2 KiB
Plaintext
110 lines
4.2 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "code",
|
|
"metadata": {},
|
|
"source": [
|
|
"# code by Tae Hwan Jung @graykode\n",
|
|
"import numpy as np\n",
|
|
"import torch\n",
|
|
"import torch.nn as nn\n",
|
|
"import torch.optim as optim\n",
|
|
"\n",
|
|
"def make_batch():\n",
|
|
" input_batch = []\n",
|
|
" target_batch = []\n",
|
|
"\n",
|
|
" words = sentence.split()\n",
|
|
" for i, word in enumerate(words[:-1]):\n",
|
|
" input = [word_dict[n] for n in words[:(i + 1)]]\n",
|
|
" input = input + [0] * (max_len - len(input))\n",
|
|
" target = word_dict[words[i + 1]]\n",
|
|
" input_batch.append(np.eye(n_class)[input])\n",
|
|
" target_batch.append(target)\n",
|
|
"\n",
|
|
" return input_batch, target_batch\n",
|
|
"\n",
|
|
"class BiLSTM(nn.Module):\n",
|
|
" def __init__(self):\n",
|
|
" super(BiLSTM, self).__init__()\n",
|
|
"\n",
|
|
" self.lstm = nn.LSTM(input_size=n_class, hidden_size=n_hidden, bidirectional=True)\n",
|
|
" self.W = nn.Linear(n_hidden * 2, n_class, bias=False)\n",
|
|
" self.b = nn.Parameter(torch.ones([n_class]))\n",
|
|
"\n",
|
|
" def forward(self, X):\n",
|
|
" input = X.transpose(0, 1) # input : [n_step, batch_size, n_class]\n",
|
|
"\n",
|
|
" hidden_state = torch.zeros(1*2, len(X), n_hidden) # [num_layers(=1) * num_directions(=2), batch_size, n_hidden]\n",
|
|
" cell_state = torch.zeros(1*2, len(X), n_hidden) # [num_layers(=1) * num_directions(=2), batch_size, n_hidden]\n",
|
|
"\n",
|
|
" outputs, (_, _) = self.lstm(input, (hidden_state, cell_state))\n",
|
|
" outputs = outputs[-1] # [batch_size, n_hidden]\n",
|
|
" model = self.W(outputs) + self.b # model : [batch_size, n_class]\n",
|
|
" return model\n",
|
|
"\n",
|
|
"if __name__ == '__main__':\n",
|
|
" n_hidden = 5 # number of hidden units in one cell\n",
|
|
"\n",
|
|
" sentence = (\n",
|
|
" 'Lorem ipsum dolor sit amet consectetur adipisicing elit '\n",
|
|
" 'sed do eiusmod tempor incididunt ut labore et dolore magna '\n",
|
|
" 'aliqua Ut enim ad minim veniam quis nostrud exercitation'\n",
|
|
" )\n",
|
|
"\n",
|
|
" word_dict = {w: i for i, w in enumerate(list(set(sentence.split())))}\n",
|
|
" number_dict = {i: w for i, w in enumerate(list(set(sentence.split())))}\n",
|
|
" n_class = len(word_dict)\n",
|
|
" max_len = len(sentence.split())\n",
|
|
"\n",
|
|
" model = BiLSTM()\n",
|
|
"\n",
|
|
" criterion = nn.CrossEntropyLoss()\n",
|
|
" optimizer = optim.Adam(model.parameters(), lr=0.001)\n",
|
|
"\n",
|
|
" input_batch, target_batch = make_batch()\n",
|
|
" input_batch = torch.FloatTensor(input_batch)\n",
|
|
" target_batch = torch.LongTensor(target_batch)\n",
|
|
"\n",
|
|
" # Training\n",
|
|
" for epoch in range(10000):\n",
|
|
" optimizer.zero_grad()\n",
|
|
" output = model(input_batch)\n",
|
|
" loss = criterion(output, target_batch)\n",
|
|
" if (epoch + 1) % 1000 == 0:\n",
|
|
" print('Epoch:', '%04d' % (epoch + 1), 'cost =', '{:.6f}'.format(loss))\n",
|
|
"\n",
|
|
" loss.backward()\n",
|
|
" optimizer.step()\n",
|
|
"\n",
|
|
" predict = model(input_batch).data.max(1, keepdim=True)[1]\n",
|
|
" print(sentence)\n",
|
|
" print([number_dict[n.item()] for n in predict.squeeze()])\n"
|
|
],
|
|
"outputs": [],
|
|
"execution_count": null
|
|
}
|
|
],
|
|
"metadata": {
|
|
"anaconda-cloud": {},
|
|
"kernelspec": {
|
|
"display_name": "Python 3",
|
|
"language": "python",
|
|
"name": "python3"
|
|
},
|
|
"language_info": {
|
|
"codemirror_mode": {
|
|
"name": "ipython",
|
|
"version": 3
|
|
},
|
|
"file_extension": ".py",
|
|
"mimetype": "text/x-python",
|
|
"name": "python",
|
|
"nbconvert_exporter": "python",
|
|
"pygments_lexer": "ipython3",
|
|
"version": "3.6.1"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 4
|
|
} |