def main():
# lstm=LayeredLSTM()
lstm = LSTM(input_size=47764,
hidden_size=512,
num_layers=8,
batch_first=True)
# this has no new sequence reset
# I wonder if gradient information will increase indefinitely
# Even so, I think detaching at the beginning of each new sequence is an arbitrary decision.
optim = torch.optim.Adam(lstm.parameters())
lstm.cuda()
criterion = torch.nn.SmoothL1Loss()
cm = ChannelManager()
cm.add_channels(2)
cm.cat_call("init_states")
for i in range(2000):
print(i)
optim.zero_grad()
target = Variable(torch.rand(2, 1, 512)).cuda()
output, states = lstm(cm.cat_call("get_input"),
cm.cat_call("get_states", 1))
cm.distribute_call("push_states", states, dim=1)
loss = criterion(output, target)
loss.backward()
optim.step()
if i % 3 == 0:
cm[0].new_sequence_reset()
if i % 5 == 0:
cm[1].new_sequence_reset()
class RecurrentEmbedding(CudaModule):
def __init__(self, input_size, hidden_size):
super(RecurrentEmbedding, self).__init__()
self.lstm = LSTM(input_size, hidden_size, 1, True, False, 0, False)
self.hidden_size = hidden_size
self.last_h = None
self.last_c = None
self.hidden_size = hidden_size
self.reset()
self.dbg_t = None
self.seq = 0
def init_weights(self):
pass
def reset(self):
self.last_h = cuda_var(torch.zeros(1, 1, self.hidden_size),
self.is_cuda, self.cuda_device)
self.last_c = cuda_var(torch.zeros(1, 1, self.hidden_size),
self.is_cuda, self.cuda_device)
def cuda(self, device=None):
CudaModule.cuda(self, device)
self.lstm.cuda(device)
return self
def forward(self, inputs):
outputs = self.lstm(inputs, (self.last_h, self.last_c))
self.last_h = outputs[1][0]
self.last_c = outputs[1][1]
return outputs[0]
class RNNOfList(StructuredEmbedding):
def __init__(self, embedding_size, hidden_size, num_layers=1,use_cuda=None):
super().__init__(hidden_size,use_cuda)
self.num_layers = num_layers
self.hidden_size = hidden_size
self.lstm = LSTM(embedding_size, hidden_size, num_layers, batch_first=True)
if use_cuda:
self.lstm=self.lstm.cuda()
def forward(self, list_of_embeddings, cuda=None):
batch_size = list_of_embeddings.size(0)
num_layers = self.num_layers
hidden_size = self.hidden_size
hidden = Variable(torch.zeros(num_layers, batch_size, hidden_size))
memory = Variable(torch.zeros(num_layers, batch_size, hidden_size))
if self.use_cuda:
hidden = hidden.cuda(async=True)
memory = memory.cuda(async=True)
states = (hidden, memory)
inputs = list_of_embeddings.view(batch_size, 1, -1)
out, states = self.lstm(inputs, states)
# return the last output:
all_outputs = out.squeeze()
if (all_outputs.dim() > 1):
# return the last output:
return all_outputs[-1, :].view(1, -1)
else:
return all_outputs.view(1, -1)
def main0():
# lstm=LayeredLSTM()
lstm = LSTM(input_size=47764,
hidden_size=128,
num_layers=8,
batch_first=True)
# this has no new sequence reset
# I wonder if gradient information will increase indefinitely
# Even so, I think detaching at the beginning of each new sequence is an arbitrary decision.
optim = torch.optim.Adam(lstm.parameters())
lstm.cuda()
h0 = Variable(torch.rand(8, 2, 128)).cuda()
c0 = Variable(torch.rand(8, 2, 128)).cuda()
states = (h0, c0)
savedetach = [states]
sd(savedetach)
for m in range(1000):
for _ in range(10):
print(_)
optim.zero_grad()
input = Variable(torch.rand(2, 1, 47764)).cuda()
target = Variable(torch.rand(2, 1, 128)).cuda()
output, states = lstm(input, savedetach[-1])
savedetach.append(states)
sd(savedetach)
criterion = torch.nn.SmoothL1Loss()
loss = criterion(output, target)
loss.backward()
optim.step()
for i in range(len(savedetach)):
del savedetach[0]
h0 = Variable(torch.rand(8, 128, 128)).cuda()
c0 = Variable(torch.rand(8, 128, 128)).cuda()
states = (h0, c0)
savedetach += [states]
class TestEncoderBase(AllenNlpTestCase):
def setUp(self):
super().setUp()
self.lstm = LSTM(bidirectional=True,
num_layers=3,
input_size=3,
hidden_size=7,
batch_first=True)
self.rnn = RNN(bidirectional=True,
num_layers=3,
input_size=3,
hidden_size=7,
batch_first=True)
self.encoder_base = _EncoderBase(stateful=True)
tensor = torch.rand([5, 7, 3])
tensor[1, 6:, :] = 0
tensor[3, 2:, :] = 0
self.tensor = tensor
mask = torch.ones(5, 7).bool()
mask[1, 6:] = False
mask[2, :] = False # <= completely masked
mask[3, 2:] = False
mask[4, :] = False # <= completely masked
self.mask = mask
self.batch_size = 5
self.num_valid = 3
sequence_lengths = get_lengths_from_binary_sequence_mask(mask)
_, _, restoration_indices, sorting_indices = sort_batch_by_length(
tensor, sequence_lengths)
self.sorting_indices = sorting_indices
self.restoration_indices = restoration_indices
def test_non_stateful_states_are_sorted_correctly(self):
encoder_base = _EncoderBase(stateful=False)
initial_states = (torch.randn(6, 5, 7), torch.randn(6, 5, 7))
# Check that we sort the state for non-stateful encoders. To test
# we'll just use a "pass through" encoder, as we aren't actually testing
# the functionality of the encoder here anyway.
_, states, restoration_indices = encoder_base.sort_and_run_forward(
lambda *x: x, self.tensor, self.mask, initial_states)
# Our input tensor had 2 zero length sequences, so we need
# to concat a tensor of shape
# (num_layers * num_directions, batch_size - num_valid, hidden_dim),
# to the output before unsorting it.
zeros = torch.zeros([6, 2, 7])
# sort_and_run_forward strips fully-padded instances from the batch;
# in order to use the restoration_indices we need to add back the two
# that got stripped. What we get back should match what we started with.
for state, original in zip(states, initial_states):
assert list(state.size()) == [6, 3, 7]
state_with_zeros = torch.cat([state, zeros], 1)
unsorted_state = state_with_zeros.index_select(
1, restoration_indices)
for index in [0, 1, 3]:
numpy.testing.assert_array_equal(
unsorted_state[:, index, :].data.numpy(),
original[:, index, :].data.numpy())
def test_get_initial_states(self):
# First time we call it, there should be no state, so we should return None.
assert (self.encoder_base._get_initial_states(
self.batch_size, self.num_valid, self.sorting_indices) is None)
# First test the case that the previous state is _smaller_ than the current state input.
initial_states = (torch.randn([1, 3, 7]), torch.randn([1, 3, 7]))
self.encoder_base._states = initial_states
# sorting indices are: [0, 1, 3, 2, 4]
returned_states = self.encoder_base._get_initial_states(
self.batch_size, self.num_valid, self.sorting_indices)
correct_expanded_states = [
torch.cat([state, torch.zeros([1, 2, 7])], 1)
for state in initial_states
]
# State should have been expanded with zeros to have shape (1, batch_size, hidden_size).
numpy.testing.assert_array_equal(
self.encoder_base._states[0].data.numpy(),
correct_expanded_states[0].data.numpy())
numpy.testing.assert_array_equal(
self.encoder_base._states[1].data.numpy(),
correct_expanded_states[1].data.numpy())
# The returned states should be of shape (1, num_valid, hidden_size) and
# they also should have been sorted with respect to the indices.
# sorting indices are: [0, 1, 3, 2, 4]
correct_returned_states = [
state.index_select(1, self.sorting_indices)[:, :self.num_valid, :]
for state in correct_expanded_states
]
numpy.testing.assert_array_equal(
returned_states[0].data.numpy(),
correct_returned_states[0].data.numpy())
numpy.testing.assert_array_equal(
returned_states[1].data.numpy(),
correct_returned_states[1].data.numpy())
# Now test the case that the previous state is larger:
original_states = (torch.randn([1, 10, 7]), torch.randn([1, 10, 7]))
self.encoder_base._states = original_states
# sorting indices are: [0, 1, 3, 2, 4]
returned_states = self.encoder_base._get_initial_states(
self.batch_size, self.num_valid, self.sorting_indices)
# State should not have changed, as they were larger
# than the batch size of the requested states.
numpy.testing.assert_array_equal(
self.encoder_base._states[0].data.numpy(),
original_states[0].data.numpy())
numpy.testing.assert_array_equal(
self.encoder_base._states[1].data.numpy(),
original_states[1].data.numpy())
# The returned states should be of shape (1, num_valid, hidden_size) and they
# also should have been sorted with respect to the indices.
correct_returned_state = [
x.index_select(1, self.sorting_indices)[:, :self.num_valid, :]
for x in original_states
]
numpy.testing.assert_array_equal(
returned_states[0].data.numpy(),
correct_returned_state[0].data.numpy())
numpy.testing.assert_array_equal(
returned_states[1].data.numpy(),
correct_returned_state[1].data.numpy())
def test_update_states(self):
assert self.encoder_base._states is None
initial_states = torch.randn([1, 5, 7]), torch.randn([1, 5, 7])
index_selected_initial_states = (
initial_states[0].index_select(1, self.restoration_indices),
initial_states[1].index_select(1, self.restoration_indices),
)
self.encoder_base._update_states(initial_states,
self.restoration_indices)
# State was None, so the updated state should just be the sorted given state.
numpy.testing.assert_array_equal(
self.encoder_base._states[0].data.numpy(),
index_selected_initial_states[0].data.numpy())
numpy.testing.assert_array_equal(
self.encoder_base._states[1].data.numpy(),
index_selected_initial_states[1].data.numpy())
new_states = torch.randn([1, 5, 7]), torch.randn([1, 5, 7])
# tensor has 2 completely masked rows, so the last 2 rows of the _sorted_ states
# will be completely zero, having been appended after calling the respective encoder.
new_states[0][:, -2:, :] = 0
new_states[1][:, -2:, :] = 0
index_selected_new_states = (
new_states[0].index_select(1, self.restoration_indices),
new_states[1].index_select(1, self.restoration_indices),
)
self.encoder_base._update_states(new_states, self.restoration_indices)
# Check that the update _preserved_ the state for the rows which were
# completely masked (2 and 4):
for index in [2, 4]:
numpy.testing.assert_array_equal(
self.encoder_base._states[0][:, index, :].data.numpy(),
index_selected_initial_states[0][:, index, :].data.numpy(),
)
numpy.testing.assert_array_equal(
self.encoder_base._states[1][:, index, :].data.numpy(),
index_selected_initial_states[1][:, index, :].data.numpy(),
)
# Now the states which were updated:
for index in [0, 1, 3]:
numpy.testing.assert_array_equal(
self.encoder_base._states[0][:, index, :].data.numpy(),
index_selected_new_states[0][:, index, :].data.numpy(),
)
numpy.testing.assert_array_equal(
self.encoder_base._states[1][:, index, :].data.numpy(),
index_selected_new_states[1][:, index, :].data.numpy(),
)
# Now test the case that the new state is smaller:
small_new_states = torch.randn([1, 3, 7]), torch.randn([1, 3, 7])
# pretend the 2nd sequence in the batch was fully masked.
small_restoration_indices = torch.LongTensor([2, 0, 1])
small_new_states[0][:, 0, :] = 0
small_new_states[1][:, 0, :] = 0
index_selected_small_states = (
small_new_states[0].index_select(1, small_restoration_indices),
small_new_states[1].index_select(1, small_restoration_indices),
)
self.encoder_base._update_states(small_new_states,
small_restoration_indices)
# Check the index for the row we didn't update is the same as the previous step:
for index in [1, 3]:
numpy.testing.assert_array_equal(
self.encoder_base._states[0][:, index, :].data.numpy(),
index_selected_new_states[0][:, index, :].data.numpy(),
)
numpy.testing.assert_array_equal(
self.encoder_base._states[1][:, index, :].data.numpy(),
index_selected_new_states[1][:, index, :].data.numpy(),
)
# Indices we did update:
for index in [0, 2]:
numpy.testing.assert_array_equal(
self.encoder_base._states[0][:, index, :].data.numpy(),
index_selected_small_states[0][:, index, :].data.numpy(),
)
numpy.testing.assert_array_equal(
self.encoder_base._states[1][:, index, :].data.numpy(),
index_selected_small_states[1][:, index, :].data.numpy(),
)
# We didn't update index 4 in the previous step either, so it should be equal to the
# 4th index of initial states.
numpy.testing.assert_array_equal(
self.encoder_base._states[0][:, 4, :].data.numpy(),
index_selected_initial_states[0][:, 4, :].data.numpy(),
)
numpy.testing.assert_array_equal(
self.encoder_base._states[1][:, 4, :].data.numpy(),
index_selected_initial_states[1][:, 4, :].data.numpy(),
)
def test_reset_states(self):
# Initialize the encoder states.
assert self.encoder_base._states is None
initial_states = torch.randn([1, 5, 7]), torch.randn([1, 5, 7])
index_selected_initial_states = (
initial_states[0].index_select(1, self.restoration_indices),
initial_states[1].index_select(1, self.restoration_indices),
)
self.encoder_base._update_states(initial_states,
self.restoration_indices)
# Check that only some of the states are reset when a mask is provided.
mask = torch.tensor([True, True, False, False, False])
self.encoder_base.reset_states(mask)
# First two states should be zeros
numpy.testing.assert_array_equal(
self.encoder_base._states[0][:, :2, :].data.numpy(),
torch.zeros_like(initial_states[0])[:, :2, :].data.numpy(),
)
numpy.testing.assert_array_equal(
self.encoder_base._states[1][:, :2, :].data.numpy(),
torch.zeros_like(initial_states[1])[:, :2, :].data.numpy(),
)
# Remaining states should be the same
numpy.testing.assert_array_equal(
self.encoder_base._states[0][:, 2:, :].data.numpy(),
index_selected_initial_states[0][:, 2:, :].data.numpy(),
)
numpy.testing.assert_array_equal(
self.encoder_base._states[1][:, 2:, :].data.numpy(),
index_selected_initial_states[1][:, 2:, :].data.numpy(),
)
# Check that error is raised if mask has wrong batch size.
bad_mask = torch.tensor([True, True, False])
with self.assertRaises(ValueError):
self.encoder_base.reset_states(bad_mask)
# Check that states are reset to None if no mask is provided.
self.encoder_base.reset_states()
assert self.encoder_base._states is None
def test_non_contiguous_initial_states_handled(self):
# Check that the encoder is robust to non-contiguous initial states.
# Case 1: Encoder is not stateful
# A transposition will make the tensors non-contiguous, start them off at the wrong shape
# and transpose them into the right shape.
encoder_base = _EncoderBase(stateful=False)
initial_states = (
torch.randn(5, 6, 7).permute(1, 0, 2),
torch.randn(5, 6, 7).permute(1, 0, 2),
)
assert not initial_states[0].is_contiguous(
) and not initial_states[1].is_contiguous()
assert initial_states[0].size() == torch.Size([6, 5, 7])
assert initial_states[1].size() == torch.Size([6, 5, 7])
# We'll pass them through an LSTM encoder and a vanilla RNN encoder to make sure it works
# whether the initial states are a tuple of tensors or just a single tensor.
encoder_base.sort_and_run_forward(self.lstm, self.tensor, self.mask,
initial_states)
encoder_base.sort_and_run_forward(self.rnn, self.tensor, self.mask,
initial_states[0])
# Case 2: Encoder is stateful
# For stateful encoders, the initial state may be non-contiguous if its state was
# previously updated with non-contiguous tensors. As in the non-stateful tests, we check
# that the encoder still works on initial states for RNNs and LSTMs.
final_states = initial_states
# Check LSTM
encoder_base = _EncoderBase(stateful=True)
encoder_base._update_states(final_states, self.restoration_indices)
encoder_base.sort_and_run_forward(self.lstm, self.tensor, self.mask)
# Check RNN
encoder_base.reset_states()
encoder_base._update_states([final_states[0]],
self.restoration_indices)
encoder_base.sort_and_run_forward(self.rnn, self.tensor, self.mask)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires cuda")
def test_non_contiguous_initial_states_handled_on_gpu(self):
# Some PyTorch operations which produce contiguous tensors on the CPU produce
# non-contiguous tensors on the GPU (e.g. forward pass of an RNN when batch_first=True).
# Accordingly, we perform the same checks from previous test on the GPU to ensure the
# encoder is not affected by which device it is on.
# Case 1: Encoder is not stateful
# A transposition will make the tensors non-contiguous, start them off at the wrong shape
# and transpose them into the right shape.
encoder_base = _EncoderBase(stateful=False).cuda()
initial_states = (
torch.randn(5, 6, 7).cuda().permute(1, 0, 2),
torch.randn(5, 6, 7).cuda().permute(1, 0, 2),
)
assert not initial_states[0].is_contiguous(
) and not initial_states[1].is_contiguous()
assert initial_states[0].size() == torch.Size([6, 5, 7])
assert initial_states[1].size() == torch.Size([6, 5, 7])
# We'll pass them through an LSTM encoder and a vanilla RNN encoder to make sure it works
# whether the initial states are a tuple of tensors or just a single tensor.
encoder_base.sort_and_run_forward(self.lstm.cuda(), self.tensor.cuda(),
self.mask.cuda(), initial_states)
encoder_base.sort_and_run_forward(self.rnn.cuda(), self.tensor.cuda(),
self.mask.cuda(), initial_states[0])
# Case 2: Encoder is stateful
# For stateful encoders, the initial state may be non-contiguous if its state was
# previously updated with non-contiguous tensors. As in the non-stateful tests, we check
# that the encoder still works on initial states for RNNs and LSTMs.
final_states = initial_states
# Check LSTM
encoder_base = _EncoderBase(stateful=True).cuda()
encoder_base._update_states(final_states,
self.restoration_indices.cuda())
encoder_base.sort_and_run_forward(self.lstm.cuda(), self.tensor.cuda(),
self.mask.cuda())
# Check RNN
encoder_base.reset_states()
encoder_base._update_states([final_states[0]],
self.restoration_indices.cuda())
encoder_base.sort_and_run_forward(self.rnn.cuda(), self.tensor.cuda(),
self.mask.cuda())
def test_layer_training(self):
"""Test AnalogLSTM layer training."""
# pylint: disable=too-many-locals, too-many-statements
def get_parameters(model, analog_if) -> dict:
"""Returns the parameter in an dict."""
dic = {}
for name, param in model.named_parameters():
if isinstance(param, AnalogContext):
weight, bias = param.analog_tile.get_weights()
splits = name.split('.')
add_on = '_' + splits[-2].split('_')[-1] + '_l' + splits[2]
dic['weight' + add_on] = weight
if bias is not None:
dic['bias' + add_on] = bias
elif analog_if and name.endswith('bias'): # digital bias
splits = name.split('.')
add_on = '_' + splits[-2].split('_')[-1] + '_l' + splits[2]
dic['bias' + add_on] = param
else:
dic[name] = param
return dic
input_size = 4
hidden_size = 3
num_layers = 1
seq_length = 10
batch_size = 3
test_for_update = False # For debugging. Does test whether all weights are updated.
# Make dataset (just random).
y_in = randn(seq_length, batch_size, input_size)
y_out = ones(seq_length, batch_size, 1)
lstm_analog = self.get_layer(input_size=input_size,
hidden_size=hidden_size,
num_layers=num_layers,
realistic_read_write=False,
dropout=0.0)
lstm = LSTM_nn(input_size=input_size,
hidden_size=hidden_size,
num_layers=num_layers,
dropout=0.0,
bias=self.bias)
weights_org = []
# pylint: disable=protected-access
lstm_analog._apply_to_analog(
lambda lay: weights_org.append(lay.analog_tile.tile.get_weights()))
lstm_pars0 = get_parameters(lstm, False)
lstm_analog_pars0 = get_parameters(lstm_analog, True)
for par_name, par_item in lstm_pars0.items():
par_item.data = lstm_analog_pars0[par_name].detach().clone()
# Make independent for comparison below.
lstm_pars0 = {
key: value.detach().clone()
for key, value in lstm_pars0.items()
}
if self.use_cuda:
y_in = y_in.cuda()
y_out = y_out.cuda()
lstm_analog.cuda()
lstm.cuda()
# First train analog and make sure weights differ.
pred_analog = self.train_once(lstm_analog,
y_in,
y_out,
True,
use_cuda=self.use_cuda)
analog_weights = []
lstm_analog._apply_to_analog(lambda lay: analog_weights.append(
lay.analog_tile.tile.get_weights()))
if test_for_update:
for weight, weight_org in zip(analog_weights, weights_org):
assert_raises(AssertionError, assert_array_almost_equal,
weight, weight_org)
# Compare with LSTM.
pred = self.train_once(lstm,
y_in,
y_out,
False,
use_cuda=self.use_cuda)
assert_array_almost_equal(pred, pred_analog)
lstm_analog._apply_to_analog(lambda lay: lay._sync_weights_from_tile())
lstm_pars = get_parameters(lstm, False)
lstm_analog_pars = get_parameters(lstm_analog, True)
if test_for_update:
for par_name, par_item in lstm_pars.items():
par0 = lstm_pars0[par_name].detach().cpu().numpy()
par = par_item.detach().cpu().numpy()
assert_raises(AssertionError, assert_array_almost_equal, par,
par0)
for par_name, par_item in lstm_pars.items():
assert_array_almost_equal(
par_item.detach().cpu().numpy(),
lstm_analog_pars[par_name].detach().cpu().numpy())