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)