import torch
import torch.nn as nn
sys.path.append("../d2l_func/")
from data_prepare import load_data_jay_song, data_iter_random, data_iter_consecutive, to_onehot
from model_train import train_rnn_pytorch
from predict import predict_rnn_pytorch
from rnn_model import RNNModel
if __name__ == "__main__":
# load data
corpus_index, char_to_idx, vocab_set, vocab_size = load_data_jay_song()
# model
hidden_num = 256
rnn_layer = nn.LSTM(vocab_size, hidden_num)
model = RNNModel(rnn_layer, vocab_size)
model = model.cuda()
loss = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
params = {
"epoch_num": 10,
"model": model,
"loss": loss,
"optimizer": optimizer,
"batch_size": 64,
"num_step": 32,
"corpus_index": corpus_index,
"data_iter": data_iter_consecutive,
"char_to_idx": char_to_idx,
"vocab_set": vocab_set,
"vocab_size": vocab_size,
total_chunks = 20 # TODO move to program args.
# Load the first chunk to get number of input features.
X_train, Y_train = tensor_loader.load_X_Y_rnn(logger, args.train_table_name, chunk=0, total_chunks=total_chunks, no_gpu=args.no_gpu)
X_val, Y_val = tensor_loader.load_X_Y_rnn(logger, args.val_table_name, chunk=0, total_chunks=total_chunks, no_gpu=args.no_gpu, validation_set=True)
N, seq_length, D_in = X_train.shape # Number of samples, sequence length, number of features.
if args.top100_labels: # Dimension of the hidden units, and dimension of the output vector.
H, D_out = 1000, 100
else:
H, D_out = 100, 10
model = RNNModel(D_in, H, D_out)
if not args.no_gpu:
model.cuda()
loss_fn = torch.nn.BCEWithLogitsLoss(size_average=True)
learning_rate, decay, momentum = 0.01, 1e-6, 0.9
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate, weight_decay=decay, momentum=momentum, nesterov=True)
tb_logger_train = tensorboardX.SummaryWriter(log_dir='../tensorboard_logs/rnn_train_' + str(experiment_id))
tb_logger_val = tensorboardX.SummaryWriter(log_dir='../tensorboard_logs/rnn_val_' + str(experiment_id))
metrics_train = defaultdict(list)
metrics_val = defaultdict(list)
metrics_test = defaultdict(list)
epochs = 3 # TODO move to program args
for chunk in range(total_chunks):
if chunk > 0: # Load next chunk (first chunk will already be loaded).
X_train, Y_train = tensor_loader.load_X_Y_rnn(logger, args.train_table_name, chunk=chunk, total_chunks=total_chunks, no_gpu=args.no_gpu)