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()
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 SimpleDPLSTMTest(unittest.TestCase):
def setUp(self):
self.SEQ_LENGTH = 20
self.INPUT_DIM = 25
self.MINIBATCH_SIZE = 30
self.LSTM_OUT_DIM = 12
self.NUM_LAYERS = 1
self.bidirectional = False
self.batch_first = False
self.num_directions = 2 if self.bidirectional else 1
self.h_init = torch.randn(
self.NUM_LAYERS * self.num_directions,
self.MINIBATCH_SIZE,
self.LSTM_OUT_DIM,
)
self.c_init = torch.randn(
self.NUM_LAYERS * self.num_directions,
self.MINIBATCH_SIZE,
self.LSTM_OUT_DIM,
)
self.original_lstm = LSTM(
self.INPUT_DIM,
self.LSTM_OUT_DIM,
batch_first=self.batch_first,
num_layers=self.NUM_LAYERS,
bidirectional=self.bidirectional,
)
self.dp_lstm = DPLSTM(
self.INPUT_DIM,
self.LSTM_OUT_DIM,
batch_first=self.batch_first,
num_layers=self.NUM_LAYERS,
bidirectional=self.bidirectional,
)
self.dp_lstm.load_state_dict(self.original_lstm.state_dict())
def _reset_seeds(self):
torch.manual_seed(1337)
torch.cuda.manual_seed(1337)
def test_lstm_forward(self):
x = (
torch.randn(self.MINIBATCH_SIZE, self.SEQ_LENGTH, self.INPUT_DIM)
if self.batch_first
else torch.randn(self.SEQ_LENGTH, self.MINIBATCH_SIZE, self.INPUT_DIM)
)
hidden = (self.h_init, self.c_init)
out, (hn, cn) = self.original_lstm(x, hidden)
dp_out, (dp_hn, dp_cn) = self.dp_lstm(x, hidden)
outputs_to_test = [
(out, dp_out, "LSTM and DPLSTM output"),
(hn, dp_hn, "LSTM and DPLSTM state `h`"),
(cn, dp_cn, "LSTM and DPLSTM state `c`"),
]
for output, dp_output, message in outputs_to_test:
assert_allclose(
actual=dp_output.expand_as(output),
expected=output,
atol=10e-6,
rtol=10e-5,
msg=f"Tensor value mismatch between {message}",
)
def test_lstm_backward(self):
x = (
torch.randn(self.MINIBATCH_SIZE, self.SEQ_LENGTH, self.INPUT_DIM)
if self.batch_first
else torch.randn(self.SEQ_LENGTH, self.MINIBATCH_SIZE, self.INPUT_DIM)
)
criterion = nn.MSELoss()
hidden = (self.h_init, self.c_init)
out, (hn, cn) = self.original_lstm(x, hidden)
y = torch.zeros_like(out)
loss = criterion(out, y)
loss.backward()
dp_out, (dp_hn, dp_cn) = self.dp_lstm(x, hidden)
dp_loss = criterion(dp_out, y)
dp_loss.backward()
dp_lstm_params = dict(self.dp_lstm.named_parameters())
for param_name, param in self.original_lstm.named_parameters():
dp_param = dp_lstm_params[param_name]
assert_allclose(
actual=dp_param,
expected=param,
atol=10e-5,
rtol=10e-3,
msg=f"Tensor value mismatch in the parameter '{param_name}'",
)
assert_allclose(
actual=dp_param.grad,
expected=param.grad,
atol=10e-6,
rtol=10e-5,
msg=f"Tensor value mismatch in the gradient of parameter '{param_name}'",
)
def test_lstm_param_update(self):
x = (
torch.randn(self.MINIBATCH_SIZE, self.SEQ_LENGTH, self.INPUT_DIM)
if self.batch_first
else torch.randn(self.SEQ_LENGTH, self.MINIBATCH_SIZE, self.INPUT_DIM)
)
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(self.original_lstm.parameters(), lr=0.5)
dp_optimizer = torch.optim.SGD(self.dp_lstm.parameters(), lr=0.5)
# Train original LSTM for one step
logits, (h_n, c_n) = self.original_lstm(x)
y = torch.zeros_like(logits)
loss = criterion(logits, y)
loss.backward()
optimizer.step()
# Train DP LSTM for one step
dp_logits, (dp_h_n, dp_c_n) = self.dp_lstm(x)
dp_loss = criterion(dp_logits, y)
dp_loss.backward()
dp_optimizer.step()
dp_lstm_params = dict(self.dp_lstm.named_parameters())
for param_name, param in self.original_lstm.named_parameters():
dp_param = dp_lstm_params[param_name]
assert_allclose(
actual=dp_param,
expected=param,
atol=10e-6,
rtol=10e-5,
msg=f"Tensor value mismatch in the parameter '{param_name}'",
)
assert_allclose(
actual=dp_param.grad,
expected=param.grad,
atol=10e-6,
rtol=10e-5,
msg=f"Tensor value mismatch in the gradient of parameter '{param_name}'",
)