307 lines
10 KiB
Python
307 lines
10 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import torch.nn as nn
|
|
import torch as T
|
|
from torch.autograd import Variable as var
|
|
import numpy as np
|
|
|
|
from torch.nn.utils.rnn import pad_packed_sequence as pad
|
|
from torch.nn.utils.rnn import pack_padded_sequence as pack
|
|
from torch.nn.utils.rnn import PackedSequence
|
|
|
|
from .util import *
|
|
from .memory import *
|
|
|
|
|
|
class DNC(nn.Module):
|
|
|
|
def __init__(
|
|
self,
|
|
input_size,
|
|
hidden_size,
|
|
rnn_type='lstm',
|
|
num_layers=1,
|
|
num_hidden_layers=2,
|
|
bias=True,
|
|
batch_first=True,
|
|
dropout=0,
|
|
bidirectional=False,
|
|
nr_cells=5,
|
|
read_heads=2,
|
|
cell_size=10,
|
|
nonlinearity='tanh',
|
|
gpu_id=-1,
|
|
independent_linears=False,
|
|
share_memory=True,
|
|
debug=False
|
|
):
|
|
super(DNC, self).__init__()
|
|
# todo: separate weights and RNNs for the interface and output vectors
|
|
|
|
self.input_size = input_size
|
|
self.hidden_size = hidden_size
|
|
self.rnn_type = rnn_type
|
|
self.num_layers = num_layers
|
|
self.num_hidden_layers = num_hidden_layers
|
|
self.bias = bias
|
|
self.batch_first = batch_first
|
|
self.dropout = dropout
|
|
self.bidirectional = bidirectional
|
|
self.nr_cells = nr_cells
|
|
self.read_heads = read_heads
|
|
self.cell_size = cell_size
|
|
self.nonlinearity = nonlinearity
|
|
self.gpu_id = gpu_id
|
|
self.independent_linears = independent_linears
|
|
self.share_memory = share_memory
|
|
self.debug = debug
|
|
|
|
self.w = self.cell_size
|
|
self.r = self.read_heads
|
|
|
|
# input size
|
|
self.nn_input_size = self.r * self.w + self.input_size
|
|
self.nn_output_size = self.r * self.w + self.hidden_size
|
|
|
|
self.interface_size = (self.w * self.r) + (3 * self.w) + (5 * self.r) + 3
|
|
self.output_size = self.hidden_size
|
|
|
|
self.rnns = [[None] * self.num_hidden_layers] * self.num_layers
|
|
self.memories = []
|
|
|
|
for layer in range(self.num_layers):
|
|
# controllers for each layer
|
|
for hlayer in range(self.num_hidden_layers):
|
|
if self.rnn_type.lower() == 'rnn':
|
|
if hlayer == 0:
|
|
self.rnns[layer][hlayer] = nn.RNNCell(self.nn_input_size, self.output_size,bias=self.bias, nonlinearity=self.nonlinearity)
|
|
else:
|
|
self.rnns[layer][hlayer] = nn.RNNCell(self.output_size, self.output_size,bias=self.bias, nonlinearity=self.nonlinearity)
|
|
elif self.rnn_type.lower() == 'gru':
|
|
if hlayer == 0:
|
|
self.rnns[layer][hlayer] = nn.GRUCell(self.nn_input_size, self.output_size, bias=self.bias)
|
|
else:
|
|
self.rnns[layer][hlayer] = nn.GRUCell(self.output_size, self.output_size, bias=self.bias)
|
|
elif self.rnn_type.lower() == 'lstm':
|
|
if hlayer == 0:
|
|
self.rnns[layer][hlayer] = nn.LSTMCell(self.nn_input_size, self.output_size, bias=self.bias)
|
|
else:
|
|
self.rnns[layer][hlayer] = nn.LSTMCell(self.output_size, self.output_size, bias=self.bias)
|
|
|
|
# memories for each layer
|
|
if not self.share_memory:
|
|
self.memories.append(
|
|
Memory(
|
|
input_size=self.output_size,
|
|
mem_size=self.nr_cells,
|
|
cell_size=self.w,
|
|
read_heads=self.r,
|
|
gpu_id=self.gpu_id,
|
|
independent_linears=self.independent_linears
|
|
)
|
|
)
|
|
|
|
# only one memory shared by all layers
|
|
if self.share_memory:
|
|
self.memories.append(
|
|
Memory(
|
|
input_size=self.output_size,
|
|
mem_size=self.nr_cells,
|
|
cell_size=self.w,
|
|
read_heads=self.r,
|
|
gpu_id=self.gpu_id,
|
|
independent_linears=self.independent_linears
|
|
)
|
|
)
|
|
|
|
for layer in range(self.num_layers):
|
|
for hlayer in range(self.num_hidden_layers):
|
|
setattr(self, 'rnn_layer_' + str(layer) + '_' + str(hlayer), self.rnns[layer][hlayer])
|
|
if not self.share_memory:
|
|
setattr(self, 'rnn_layer_memory_' + str(layer), self.memories[layer])
|
|
if self.share_memory:
|
|
setattr(self, 'rnn_layer_memory_shared', self.memories[0])
|
|
|
|
# final output layer
|
|
self.output_weights = nn.Linear(self.output_size, self.output_size)
|
|
self.mem_out = nn.Linear(self.nn_output_size, self.input_size)
|
|
self.dropout_layer = nn.Dropout(self.dropout)
|
|
|
|
if self.gpu_id != -1:
|
|
[x.cuda(self.gpu_id) for y in self.rnns for x in y]
|
|
[x.cuda(self.gpu_id) for x in self.memories]
|
|
self.mem_out.cuda(self.gpu_id)
|
|
|
|
def _init_hidden(self, hx, batch_size, reset_experience):
|
|
# create empty hidden states if not provided
|
|
if hx is None:
|
|
hx = (None, None, None)
|
|
(chx, mhx, last_read) = hx
|
|
|
|
# initialize hidden state of the controller RNN
|
|
if chx is None:
|
|
chx = cuda(T.zeros(batch_size, self.output_size), gpu_id=self.gpu_id)
|
|
if self.rnn_type.lower() == 'lstm':
|
|
chx = [ [ (chx.clone(), chx.clone()) for h in range(self.num_hidden_layers) ] for l in range(self.num_layers) ]
|
|
else:
|
|
chx = [ [ chx.clone() for h in range(self.num_hidden_layers) ] for l in range(self.num_layers) ]
|
|
|
|
# Last read vectors
|
|
if last_read is None:
|
|
last_read = cuda(T.zeros(batch_size, self.w * self.r), gpu_id=self.gpu_id)
|
|
|
|
# memory states
|
|
if mhx is None:
|
|
if self.share_memory:
|
|
mhx = self.memories[0].reset(batch_size, erase=reset_experience)
|
|
else:
|
|
mhx = [m.reset(batch_size, erase=reset_experience) for m in self.memories]
|
|
else:
|
|
if self.share_memory:
|
|
mhx = self.memories[0].reset(batch_size, mhx, erase=reset_experience)
|
|
else:
|
|
mhx = [m.reset(batch_size, h, erase=reset_experience) for m, h in zip(self.memories, mhx)]
|
|
|
|
return chx, mhx, last_read
|
|
|
|
def _debug(self, mhx, debug_obj):
|
|
debug_obj['memory'].append(mhx['memory'][0].data.cpu().numpy())
|
|
debug_obj['link_matrix'].append(mhx['link_matrix'][0][0].data.cpu().numpy())
|
|
debug_obj['precedence'].append(mhx['precedence'][0].data.cpu().numpy())
|
|
debug_obj['read_weights'].append(mhx['read_weights'][0].data.cpu().numpy())
|
|
debug_obj['write_weights'].append(mhx['write_weights'][0].data.cpu().numpy())
|
|
debug_obj['usage_vector'].append(mhx['usage_vector'][0].unsqueeze(0).data.cpu().numpy())
|
|
|
|
def _layer_forward(self, input, layer, hx=(None, None)):
|
|
(chx, mhx) = hx
|
|
max_length = len(input)
|
|
outs = [0] * max_length
|
|
read_vectors = [0] * max_length
|
|
|
|
if self.debug:
|
|
mem_debug = {
|
|
"memory": [],
|
|
"link_matrix": [],
|
|
"precedence": [],
|
|
"read_weights": [],
|
|
"write_weights": [],
|
|
"usage_vector": [],
|
|
}
|
|
|
|
for time in range(max_length):
|
|
# pass through controller
|
|
layer_input = input[time]
|
|
hchx = []
|
|
|
|
for hlayer in range(self.num_hidden_layers):
|
|
h = self.rnns[layer][hlayer](layer_input, chx[hlayer])
|
|
layer_input = h[0] if self.rnn_type.lower() == 'lstm' else h
|
|
hchx.append(h)
|
|
chx = hchx
|
|
|
|
# the interface vector
|
|
ξ = layer_input
|
|
# the output
|
|
out = self.output_weights(layer_input)
|
|
|
|
# pass through memory
|
|
if self.share_memory:
|
|
read_vecs, mhx = self.memories[0](ξ, mhx)
|
|
if self.debug:
|
|
self._debug(mhx, mem_debug)
|
|
else:
|
|
read_vecs, mhx = self.memories[layer](ξ, mhx)
|
|
if self.debug:
|
|
self._debug(mhx, mem_debug)
|
|
read_vectors[time] = read_vecs.view(-1, self.w * self.r)
|
|
|
|
# get the final output for this time step
|
|
outs[time] = self.dropout_layer(self.mem_out(T.cat([out, read_vectors[time]], 1)))
|
|
|
|
if self.debug:
|
|
return outs, read_vectors, mem_debug, (chx, mhx)
|
|
else:
|
|
return outs, read_vectors, (chx, mhx)
|
|
|
|
def forward(self, input, hx=(None, None, None), reset_experience=False):
|
|
# handle packed data
|
|
is_packed = type(input) is PackedSequence
|
|
if is_packed:
|
|
input, lengths = pad(input)
|
|
max_length = lengths[0]
|
|
else:
|
|
max_length = input.size(1) if self.batch_first else input.size(0)
|
|
lengths = [input.size(1)] * max_length if self.batch_first else [input.size(0)] * max_length
|
|
|
|
batch_size = input.size(0) if self.batch_first else input.size(1)
|
|
|
|
# make the data batch-first
|
|
if not self.batch_first:
|
|
input = input.transpose(0, 1)
|
|
|
|
controller_hidden, mem_hidden, last_read = self._init_hidden(hx, batch_size, reset_experience)
|
|
|
|
# batched forward pass per element / word / etc
|
|
outputs = None
|
|
chxs = []
|
|
if self.debug:
|
|
viz = {}
|
|
|
|
read_vectors = [last_read] * max_length
|
|
# outs = [input[:, x, :] for x in range(max_length)]
|
|
outs = [T.cat([input[:, x, :], last_read], 1) for x in range(max_length)]
|
|
|
|
for layer in range(self.num_layers):
|
|
# this layer's hidden states
|
|
chx = controller_hidden[layer]
|
|
|
|
m = mem_hidden if self.share_memory else mem_hidden[layer]
|
|
# pass through controller
|
|
if self.debug:
|
|
outs, _, mem_debug, (chx, m) = self._layer_forward(outs,layer,(chx, m))
|
|
else:
|
|
outs, _, (chx, m) = self._layer_forward(outs, layer, (chx, m))
|
|
|
|
# debug memory
|
|
if self.debug:
|
|
if viz == {}:
|
|
viz = mem_debug
|
|
else:
|
|
viz["memory"] += mem_debug["memory"]
|
|
viz["link_matrix"] += mem_debug["link_matrix"]
|
|
viz["precedence"] += mem_debug["precedence"]
|
|
viz["read_weights"] += mem_debug["read_weights"]
|
|
viz["write_weights"] += mem_debug["write_weights"]
|
|
viz["usage_vector"] += mem_debug["usage_vector"]
|
|
|
|
# store the memory back (per layer or shared)
|
|
if self.share_memory:
|
|
mem_hidden = m
|
|
else:
|
|
mem_hidden[layer] = m
|
|
chxs.append(chx)
|
|
|
|
if layer == self.num_layers - 1:
|
|
# final outputs
|
|
outputs = T.stack(outs, 1)
|
|
else:
|
|
# the controller output + read vectors go into next layer
|
|
outs = [T.cat([o, r], 1) for o, r in zip(outs, read_vectors)]
|
|
# outs = [o for o in outs]
|
|
|
|
if self.debug:
|
|
viz = { k: np.array(v) for k,v in viz.items() }
|
|
viz = { k: v.reshape(v.shape[0], v.shape[1] * v.shape[2]) for k,v in viz.items() }
|
|
|
|
controller_hidden = chxs
|
|
|
|
if not self.batch_first:
|
|
outputs = outputs.transpose(0, 1)
|
|
if is_packed:
|
|
outputs = pack(output, lengths)
|
|
|
|
if self.debug:
|
|
return outputs, (controller_hidden, mem_hidden, read_vectors[-1]), viz
|
|
else:
|
|
return outputs, (controller_hidden, mem_hidden, read_vectors[-1])
|