def sample(args):
with open(os.path.join(args.save_dir, 'config.pkl'), 'rb') as f:
saved_args = cPickle.load(f)
with open(os.path.join(args.save_dir, 'chars_vocab.pkl'), 'rb') as f:
chars, vocab = cPickle.load(f)
model = RNN_Model(saved_args, mode='INFER')
with tf.Session() as sess:
tf.global_variables_initializer().run()
saver = tf.train.Saver(tf.global_variables())
ckpt = tf.train.get_checkpoint_state(args.save_dir)
if ckpt and ckpt.model_checkpoint_path:
saver.restore(sess=sess, save_path=ckpt.model_checkpoint_path)
print(
model.sample(sess, chars, vocab, args.num_samples, args.prime,
args.sample))
def train(args):
data_loader = Data_loader(args.data_dir,args.batch_size,args.seq_length)
args.vocab_size = data_loader.vocab_size
# Save configurations
with open(os.path.join(args.save_dir,'config.pkl'),'wb') as f:
cPickle.dump(args,f)
# Save (chars, vocab) pairs
with open(os.path.join(args.save_dir,'chars_vocab.pkl'),'wb') as f:
cPickle.dump((data_loader.chars,data_loader.vocab), f)
model = RNN_Model(args)
with tf.Session() as sess:
tf.global_variables_initializer().run()
saver = tf.train.Saver(tf.global_variables())
tf.summary.FileWriter(logdir=os.path.join(args.save_dir,'graph'),
graph=sess.graph).add_graph(sess.graph)
for e in range(args.num_epochs):
sess.run(tf.assign(model.lr,args.learning_rate*args.decay_rate**e))
data_loader.reset_batch_pointer()
state = sess.run(model.initial_state)
for b in range(data_loader.num_batches):
start = time.time()
x,y = data_loader.next_batch()
feed = {model.input_data:x, model.targets: y}
for i, (c,h) in enumerate(model.initial_state):
feed[c] = state[i].c
feed[h] = state[i].h
train_loss,state,_ = sess.run([model.cost,model.final_state,model.train_op],
feed_dict = feed)
end = time.time()
print('{}/{} (epoch{}), train_loss = {:.3f},time/batch={:.3f}'.format(
e*data_loader.num_batches+b,args.num_epochs*data_loader.num_batches,
e,train_loss,end-start))
# save for the last result
if (e==args.num_epochs - 1 and b==data_loader.num_batches - 1):
checkpoint_path = os.path.join(args.save_dir,'model.ckpt')
saver.save(sess,checkpoint_path,
global_step = e*data_loader.num_batches + b)
print('model save to {}'.format(checkpoint_path))
torch.manual_seed(args.seed)
data_dir = "data/" + args.data + "/"
if args.model == "main":
network = MainModel(args.d_marker, args.neighbor, args.d_model, args.d_inner, args.d_model,
args.d_model, args.d_model, args.d_head, args.max_time, args.embed_ratio,
args.cuda, args.sample, args.discount, args.regular, dropout=args.dropout).cuda(args.cuda)
if args.model == "pr":
network = PR_Model(args.d_marker, args.neighbor, args.d_model, args.d_inner, args.d_model,
args.d_model, args.d_model, args.d_head, args.max_time, args.embed_ratio,
args.cuda, args.sample, args.discount, args.regular, dropout=args.dropout).cuda(args.cuda)
if args.model == "rnn":
network = RNN_Model(args.d_marker, args.neighbor, args.d_model, args.max_time, args.embed_ratio, args.cuda,
args.sample, args.discount, args.regular, dropout=args.dropout).cuda(args.cuda)
network.generator.sample_linear.cpu()
marker_data = torch.load(data_dir + "marker.pkl").cuda(args.cuda)
time_data = torch.load(data_dir + "time.pkl").cuda(args.cuda)
mask_data = torch.load(data_dir + "mask.pkl").cuda(args.cuda)
adj_list = torch.load(data_dir + "adj_list.pkl")
train_data_num = int(args.d_cascade / 10 * 9)
train_marker = marker_data[0: train_data_num, :]
train_time = time_data[0: train_data_num, :]
train_mask = mask_data[0: train_data_num, :]
test_marker = marker_data[train_data_num:, :]
vocab_size = 4
for word in vocab:
word2id[word] = vocab_size
vocab_size += 1
region2id , numRegions = getRegionLabels(train_set)
period2id , numPeriods = getPeriodLabels(train_set)
train_input , train_output_region , train_output_time = getInput(train_set , word2id , period2id , region2id , maxL)
test_input , test_output_region , test_output_time = getInput(test_set , word2id , period2id , region2id , maxL)
print(train_set[0])
print(train_input[0])
print(train_output_region[0])
print(train_output_time[0])
model = RNN_Model(embed_size=args.embed_size,
hidden_size=args.hidden_size,
vocab_len=vocab_size,
epoch=args.epoch_size,
learning_rate=args.lr,
batch_size=args.batch_size)
model.train(numpy.asarray(train_input) , numpy.asarray(train_output_region) , numpy.asarray(train_output_time))
model.test(test_input , test_output_region , test_output_time)
print()
print()
print("Training Results:")
print()
model.test(train_input , train_output_region , train_output_time)
if __name__ == "__main__":
config = parse()
config.model_name = "{}-{}-{}-{}".format(config.data_type,
config.direction,
config.cell_type,
str(config.layers))
config.model_dir = join(config.output_dir, config.model_name)
if not exists(config.output_dir):
mkdir(config.output_dir)
if not exists(config.model_dir):
mkdir(config.model_dir)
model = RNN_Model(config)
if config.testing:
test_ids, test_data, _ = load_timit(config, data_set='test')
test_ids = np.expand_dims(test_ids, 1)
predictions = model.test(test_data)
predictions = np.expand_dims(
post_processing(config, predictions, threshold=2), 1)
outputs = np.append(test_ids, predictions, axis=1)
df = pd.DataFrame(outputs).to_csv(join(
config.model_dir, '{}.csv'.format(config.model_name)),
index=False,
header=['id', 'phone_sequence'])