from os.path import join
import torch
import torch.optim as optim
from torch.utils.tensorboard import SummaryWriter
from visualization import plot_piano_roll
# from tensorboardX import SummaryWriter
# set config
# parser = custom.get_argument_parser()
# args = parser.parse_args("-m train_model -c config/train.yml")
# config.load(args.model_dir, args.configs, initialize=True)
# model_dir = r"E:\Github_Projects\music_DeepLearning\MusicTransformer-pytorch\train_model\longer_seq"
# configs = [r"E:\Github_Projects\music_DeepLearning\MusicTransformer-pytorch\train_model\longer_seq\train.yml"]
model_dir = r"C:\Users\PonceLab\Documents\MusicTransformer-pytorch\train_model\1536_seq"
configs = [join(model_dir, r"train.yml")]
config.load(model_dir, configs, initialize=False)
# check cuda
if torch.cuda.is_available():
config.device = torch.device('cuda')
else:
config.device = torch.device('cpu')
# load data
dataset = Data(config.pickle_dir)
print(dataset)
# load model
learning_rate = config.l_r
# define model
from custom.criterion import SmoothCrossEntropyLoss, CustomSchedule
from custom.config import config
from data import Data
import utils
import datetime
import time
import torch
import torch.optim as optim
from tensorboardX import SummaryWriter
# set config
parser = custom.get_argument_parser()
args = parser.parse_args()
config.load(args.model_dir, args.configs, initialize=True)
print("args.model_dir ", args.model_dir)
print("args.configs ", args.configs)
# check cuda
if torch.cuda.is_available():
config.device = torch.device('cuda')
else:
config.device = torch.device('cpu')
# load data
dataset = Data(config.pickle_dir)
print(dataset)
# load model
learning_rate = config.l_r
from model import MusicTransformer
import custom
from custom.config import config
import torch
parser = custom.get_argument_parser()
args = parser.parse_args()
config.load(args.model_dir, [args.model_dir + '/save.yml'] + args.configs,
initialize=True)
# # check cuda
# if torch.cuda.is_available():
# config.device = torch.device('cuda')
# else:
config.device = torch.device('cpu')
mt = MusicTransformer(embedding_dim=config.embedding_dim,
vocab_size=config.vocab_size,
num_layer=config.num_layers,
max_seq=config.max_seq,
dropout=0,
debug=False)
mt.load_state_dict(torch.load(args.model_dir + '/final.pth'))
mt.test()
mt_script = torch.jit.trace(mt, (torch.rand(1, 1), torch.tensor(100)))
print(mt_script.code)