def test_lstm():
h_c = 2
active_func = nn.Tanh()
in_c = 1
in_h = 4
in_w = 4
kernel_size = 3
time_steps = 5
DEBUG = True
num_layers = 3
cell_conf_l0 = pack(
[h_c, active_func, in_c, in_h, in_w, kernel_size, DEBUG],
ConvLSTMCell.get_init_keys())
cell_conf_l1 = pack(
[h_c, active_func, h_c, in_h, in_w, kernel_size, DEBUG],
ConvLSTMCell.get_init_keys())
cell_conf_l2 = pack(
[h_c, active_func, h_c, in_h, in_w, kernel_size, DEBUG],
ConvLSTMCell.get_init_keys())
cell_configs = [cell_conf_l0, cell_conf_l1, cell_conf_l2]
configs = pack([num_layers, cell_configs], ConvLSTM.get_init_keys())
model = ConvLSTM(configs)
batch_size = 3
input = Variable(torch.randn(batch_size, time_steps, in_c, in_h, in_w))
data = pack([input, model.init_hidden(batch_size, False)], ['x', 'states'])
# print (data)
out = model(data)
print(out)
def test_lstm_cell():
h_c = 16
active_func = nn.Tanh()
in_c = 1
in_h = 32
in_w = 32
kernel_size = 3
DEBUG = True
configs = pack([h_c, active_func, in_c, in_h, in_w, kernel_size, DEBUG],
ConvLSTMCell.get_init_keys())
model = ConvLSTMCell(configs)
batch_size = 4
input = Variable(torch.randn(batch_size, in_c, in_h, in_w))
data = pack([input, model.init_hidden(batch_size, False)], ['x', 'states'])
# print (data)
out = model(data)
print(out)
def forward(self, data, configs=None):
x_train, x_predict, states = unpack(data,
['x_train', 'x_predict', 'states'])
use_gt, max_steps = unpack(configs, ['use_gt', 'max_steps'])
# x: batch_size * time_steps * channels * height * width
# states: batch_size * channels' * height * width
batch_size = x_train.size(0)
time_steps = x_train.size(1)
encoder = self.encoder
reconstructor = self.reconstructor
predictor = self.predictor
# encoding stages
en_data = pack([x_train, encoder.init_hidden(batch_size)],
['x', 'states'])
states = encoder(en_data)
# reconstruct
r_res = []
r_x = None
r_states = states
for t in xrange(time_steps):
r_data = pack([r_x, r_states], ['x', 'states'])
r_out, r_states = reconstructor(r_data)
r_res.append(r_out)
r_x = x_train[:, t].unsqueeze(1) if use_gt else r_out.unsqueeze(1)
# predict
time_steps = x_predict.size(1) if use_gt else max_steps
p_res = []
p_x = None
p_states = states
# print ('start from p')
for t in xrange(time_steps):
p_data = pack([p_x, p_states], ['x', 'states'])
p_out, p_states = predictor(p_data)
p_res.append(p_out)
p_x = x_predict[:,
t].unsqueeze(1) if use_gt else p_out.unsqueeze(1)
return torch.cat(r_res, 1), torch.cat(p_res, 1)
def val(dataloader, model, val_info, criterion):
model.eval()
num_epochs = val_info['num_epochs']
epoch = val_info['epoch']
total_steps = len(dataloader)
total_loss = 0
for idx, (x_train, x_predict) in enumerate(dataloader):
x_train = to_var(x_train)
x_predict = to_var(x_predict)
data = pack([x_train, x_predict, None],
['x_train', 'x_predict', 'states'])
configs = pack([False, 10], ['use_gt', 'max_steps'])
reconstruct, predict = model(data, configs)
r_loss = criterion(reconstruct, x_train)
p_loss = criterion(predict, x_predict)
loss = r_loss + p_loss
logger.info(
'[Val] Epoch [%d/%d], Step [%d/%d], Reconstruct Loss: %5.4f, Predict Loss: %5.4f, Total: %5.4f'
% (epoch, num_epochs, idx + 1, total_steps, r_loss.data[0],
p_loss.data[0], loss.data[0]))
total_loss += loss.data[0]
return total_loss / total_steps
def train(dataloader, model, optimizer, criterion, train_info):
model.train()
num_epochs = train_info['num_epochs']
epoch = train_info['epoch']
clip = train_info['clip']
total_steps = len(dataloader)
for idx, (x_train, x_predict) in enumerate(dataloader):
optimizer.zero_grad()
x_train = to_var(x_train)
x_predict = to_var(x_predict)
data = pack([x_train, x_predict, None],
['x_train', 'x_predict', 'states'])
configs = pack([False, 10], ['use_gt', 'max_steps'])
reconstruct, predict = model(data, configs)
r_loss = criterion(reconstruct, x_train)
p_loss = criterion(predict, x_predict)
loss = r_loss + p_loss
loss.backward()
torch.nn.utils.clip_grad_norm(model.parameters(), clip)
optimizer.step()
logger.info(
'Epoch [%d/%d], Step [%d/%d], Reconstruct Loss: %5.4f, Predict Loss: %5.4f, Total: %5.4f'
% (epoch, num_epochs, idx + 1, total_steps, r_loss.data[0],
p_loss.data[0], loss.data[0]))
def forward(self, data, configs=None):
_KEYS = ['x', 'states']
x, states = unpack(data, _KEYS)
batch_size = states[0][0].size(0)
if x is None:
x_c, x_h, x_w = self.cell_config['in_c'], self.cell_config[
'in_w'], self.cell_config['in_h']
x = to_var(torch.zeros(batch_size, 1, x_c, x_h, x_w))
# x: batch_size, time_steps, channels, height, width
time_steps = x.size(1)
next_states = []
cell_list = self.cell_list
current_input = [x[:, t] for t in xrange(time_steps)]
for l in xrange(self.num_layers):
h0, c0 = states[l]
for t in xrange(time_steps):
data = pack([current_input[t], (h0, c0)], ['x', 'states'])
h, c = cell_list[l](data)
next_states.append((h, c))
current_input[t] = h
states[l] = (h, c)
return states
if __name__ == '__main__':
_KEYS = ['encoder_configs', 'reconstruct_configs', 'predict_configs']
h_c = 2
active_func = nn.Tanh()
in_c = 1
in_h = 4
in_w = 4
kernel_size = 3
time_steps = 5
DEBUG = True
batch_size = 3
num_layers = 3
cell_conf_l0 = pack(
[h_c, active_func, in_c, in_h, in_w, kernel_size, DEBUG],
ConvLSTMCell.get_init_keys())
cell_conf_l1 = pack(
[h_c, active_func, h_c, in_h, in_w, kernel_size, DEBUG],
ConvLSTMCell.get_init_keys())
cell_conf_l2 = pack(
[h_c, active_func, h_c, in_h, in_w, kernel_size, DEBUG],
ConvLSTMCell.get_init_keys())
cell_configs = [cell_conf_l0, cell_conf_l1, cell_conf_l2]
encoder_configs = pack([num_layers, cell_configs],
ConvLSTM.get_init_keys())
reconstruct_configs = pack([num_layers, cell_configs],
ConvLSTM.get_init_keys())
predict_configs = pack([num_layers, cell_configs],
ConvLSTM.get_init_keys())
def MMNIST_CONV_LSTM(extra_info):
_KEYS = ['encoder_configs', 'reconstruct_configs', 'predict_configs']
h_c = 16
active_func = nn.Tanh()
in_c = 1
in_h = 64
in_w = 64
kernel_size = 5
DEBUG = True
num_layers = 3
cell_conf_l0 = pack(
[h_c, active_func, in_c, in_h, in_w, kernel_size, DEBUG],
ConvLSTMCell.get_init_keys())
cell_conf_l1 = pack([8, active_func, h_c, in_h, in_w, kernel_size, DEBUG],
ConvLSTMCell.get_init_keys())
cell_conf_l2 = pack([8, active_func, 8, in_h, in_w, kernel_size, DEBUG],
ConvLSTMCell.get_init_keys())
cell_configs = [cell_conf_l0, cell_conf_l1, cell_conf_l2]
encoder_configs = pack([num_layers, cell_configs],
ConvLSTM.get_init_keys())
reconstruct_configs = pack([num_layers, cell_configs],
ConvLSTM.get_init_keys())
predict_configs = pack([num_layers, cell_configs],
ConvLSTM.get_init_keys())
model_info = pack([encoder_configs, reconstruct_configs, predict_configs],
_KEYS)
model_info['name'] = 'MMNIST_CONV_LSTM'
trainloader_info = {
'file_addr': './data/mmnist_train.npy',
'batch_size': 32,
'shuffle': True,
'num_workers': 2
}
valloader_info = {
'file_addr': './data/mmnist_val.npy',
'batch_size': 16,
'shuffle': False,
'num_workers': 2
}
testloader_info = {
'file_addr': './data/mmnist_test.npy',
'batch_size': 16,
'shuffle': False,
'num_workers': 2
}
seed = 666
folder_name = 'mmnist_convLSTM'
main_info = {
'clip': 0.25,
'num_epochs': 60,
'halve_every': 10,
'log_dir': './logs/%s' % folder_name,
'save_dir': './checkpoints/%s' % folder_name
}
optimizer_info = {
'lr': 1e-4,
'optim_alg': 'RMSprop',
'weight_decay': 0.9,
'momentum': 0
}
hparams = HParams(trainloader_info=trainloader_info,
valloader_info=valloader_info,
testloader_info=testloader_info,
model_info=model_info,
optimizer_info=optimizer_info,
main_info=main_info,
seed=seed)
return hparams