def __init__(self, opt):
super(SeqGANInstructor, self).__init__(opt)
# generator, discriminator
self.gen = SeqGAN_G(cfg.gen_embed_dim,
cfg.gen_hidden_dim,
cfg.vocab_size,
cfg.max_seq_len,
cfg.padding_idx,
cfg.temperature,
gpu=cfg.CUDA)
self.dis = SeqGAN_D(cfg.dis_embed_dim,
cfg.vocab_size,
cfg.padding_idx,
gpu=cfg.CUDA)
self.init_model()
# Optimizer
self.gen_opt = optim.Adam(self.gen.parameters(), lr=cfg.gen_lr)
self.dis_opt = optim.Adam(self.dis.parameters(), lr=cfg.dis_lr)
# Criterion
self.mle_criterion = nn.NLLLoss()
self.dis_criterion = nn.CrossEntropyLoss()
# DataLoader
self.gen_data = GenDataIter(
self.gen.sample(cfg.batch_size, cfg.batch_size))
self.dis_data = DisDataIter(self.gen_data.random_batch()['target'],
self.oracle_data.random_batch()['target'])
self.dis_eval_data = DisDataIter(
self.gen_data.random_batch()['target'],
self.oracle_data.random_batch()['target'])
def __init__(self, opt):
super(LeakGANInstructor, self).__init__(opt)
# generator, discriminator
self.gen = LeakGAN_G(cfg.gen_embed_dim, cfg.gen_hidden_dim,
cfg.vocab_size, cfg.max_seq_len, cfg.padding_idx,
cfg.goal_size, cfg.step_size, cfg.CUDA)
self.dis = LeakGAN_D(cfg.dis_embed_dim,
cfg.vocab_size,
cfg.padding_idx,
gpu=cfg.CUDA)
self.init_model()
# optimizer
mana_params, work_params = self.gen.split_params()
mana_opt = optim.Adam(mana_params, lr=cfg.gen_lr)
work_opt = optim.Adam(work_params, lr=cfg.gen_lr)
self.gen_opt = [mana_opt, work_opt]
self.dis_opt = optim.Adam(self.dis.parameters(), lr=cfg.dis_lr)
# Criterion
self.mle_criterion = nn.NLLLoss()
self.dis_criterion = nn.CrossEntropyLoss()
# DataLoader
self.gen_data = GenDataIter(
self.gen.sample(cfg.batch_size, cfg.batch_size, self.dis))
def __init__(self, opt):
self.log = create_logger(__name__, silent=False, to_disk=True,
log_file=cfg.log_filename if cfg.if_test
else [cfg.log_filename, cfg.save_root + 'log.txt'])
self.sig = Signal(cfg.signal_file)
self.opt = opt
# oracle, generator, discriminator
self.oracle = Oracle(cfg.gen_embed_dim, cfg.gen_hidden_dim, cfg.vocab_size, cfg.max_seq_len,
cfg.padding_idx, gpu=cfg.CUDA)
self.oracle_list = [Oracle(cfg.gen_embed_dim, cfg.gen_hidden_dim, cfg.vocab_size, cfg.max_seq_len,
cfg.padding_idx, gpu=cfg.CUDA) for _ in range(cfg.k_label)]
self.dis = None
self.clas = None
self.show_config()
self.check_oracle() # Create Oracle models if not exist
# DataLoader
self.oracle_samples = torch.load(cfg.oracle_samples_path.format(cfg.samples_num))
self.oracle_samples_list = [torch.load(cfg.multi_oracle_samples_path.format(i, cfg.samples_num))
for i in range(cfg.k_label)]
self.oracle_data = GenDataIter(self.oracle_samples)
self.oracle_data_list = [GenDataIter(self.oracle_samples_list[i]) for i in range(cfg.k_label)]
# Criterion
self.mle_criterion = nn.NLLLoss()
self.dis_criterion = nn.CrossEntropyLoss()
# Metrics
self.nll_oracle = NLL('NLL_oracle', if_use=cfg.use_nll_oracle, gpu=cfg.CUDA)
self.nll_gen = NLL('NLL_gen', if_use=cfg.use_nll_gen, gpu=cfg.CUDA)
self.nll_div = NLL('NLL_div', if_use=cfg.use_nll_div, gpu=cfg.CUDA)
self.all_metrics = [self.nll_oracle, self.nll_gen, self.nll_div]
def __init__(self, opt):
super(LeakGANInstructor, self).__init__(opt)
# generator, discriminator
self.gen = LeakGAN_G(cfg.gen_embed_dim, cfg.gen_hidden_dim, cfg.vocab_size, cfg.max_seq_len,
cfg.padding_idx, cfg.goal_size, cfg.step_size, cfg.CUDA)
self.dis = LeakGAN_D(cfg.dis_embed_dim, cfg.vocab_size, cfg.padding_idx, gpu=cfg.CUDA)
self.init_model()
# optimizer
mana_params, work_params = self.gen.split_params()
mana_opt = optim.Adam(mana_params, lr=cfg.gen_lr)
work_opt = optim.Adam(work_params, lr=cfg.gen_lr)
self.gen_opt = [mana_opt, work_opt]
self.dis_opt = optim.Adam(self.dis.parameters(), lr=cfg.dis_lr)
# Criterion
self.mle_criterion = nn.NLLLoss()
self.dis_criterion = nn.BCEWithLogitsLoss()
# DataLoader
self.gen_data = GenDataIter(self.gen.sample(cfg.batch_size, cfg.batch_size, self.dis))
self.dis_data = DisDataIter(self.gen_data.random_batch()['target'], self.train_data.random_batch()['target'])
# Metrics
self.bleu = BLEU(test_text=tensor_to_tokens(self.gen_data.target, self.index_word_dict),
real_text=tensor_to_tokens(self.test_data.target, self.test_data.index_word_dict), gram=3)
self.self_bleu = BLEU(test_text=tensor_to_tokens(self.gen_data.target, self.index_word_dict),
real_text=tensor_to_tokens(self.gen_data.target, self.index_word_dict),
gram=3)
def __init__(self, opt):
super(SeqGANInstructor, self).__init__(opt)
# generator, discriminator
self.gen = SeqGAN_G(cfg.gen_embed_dim, cfg.gen_hidden_dim, cfg.vocab_size, cfg.max_seq_len,
cfg.padding_idx, cfg.temperature, gpu=cfg.CUDA)
self.dis = SeqGAN_D(cfg.dis_embed_dim, cfg.vocab_size, cfg.padding_idx, gpu=cfg.CUDA)
self.init_model()
# Optimizer
self.gen_opt = optim.Adam(self.gen.parameters(), lr=cfg.gen_lr)
self.gen_adv_opt = optim.Adam(self.gen.parameters(), lr=cfg.gen_lr)
self.dis_opt = optim.Adam(self.dis.parameters(), lr=cfg.dis_lr)
# Criterion
self.mle_criterion = nn.NLLLoss()
self.dis_criterion = nn.CrossEntropyLoss()
# DataLoader
self.gen_data = GenDataIter(self.gen.sample(cfg.batch_size, cfg.batch_size))
self.dis_data = DisDataIter(self.train_data.random_batch()['target'], self.gen_data.random_batch()['target'])
# Metrics
self.bleu = BLEU(test_text=tensor_to_tokens(self.gen_data.target, self.index_word_dict),
real_text=tensor_to_tokens(self.test_data.target, self.test_data.index_word_dict), gram=3)
self.self_bleu = BLEU(test_text=tensor_to_tokens(self.gen_data.target, self.index_word_dict),
real_text=tensor_to_tokens(self.gen_data.target, self.index_word_dict),
gram=3)
def __init__(self, opt):
super(RelbarGANInstructor, self).__init__(opt)
# generator, discriminator
self.gen = RelbarGAN_G(cfg.mem_slots,
cfg.num_heads,
cfg.head_size,
cfg.gen_embed_dim,
cfg.gen_hidden_dim,
cfg.vocab_size,
cfg.max_seq_len,
cfg.padding_idx,
cfg.temperature,
cfg.eta,
gpu=cfg.CUDA)
self.dis = RelbarGAN_D(cfg.dis_embed_dim,
cfg.max_seq_len,
cfg.num_rep,
cfg.vocab_size,
cfg.padding_idx,
gpu=cfg.CUDA)
self.init_model()
# Optimizer
self.gen_opt = optim.Adam(self.gen.parameters(), lr=cfg.gen_lr)
self.gen_adv_opt = optim.Adam(itertools.chain(
self.gen.parameters(), [self.gen.temperature, self.gen.eta]),
lr=cfg.gen_adv_lr)
self.dis_opt = optim.Adam(self.dis.parameters(), lr=cfg.dis_lr)
# Criterion
self.mle_criterion = nn.NLLLoss()
self.dis_pretrain_criterion = nn.BCEWithLogitsLoss()
# DataLoader
self.gen_data = GenDataIter(
self.gen.sample(cfg.batch_size, cfg.batch_size))
self.dis_data = DisDataIter(self.train_data.random_batch()['target'],
self.gen_data.random_batch()['target'])
# Metrics
bleu_gram = list(range(2, cfg.max_seq_len +
1)) if cfg.max_seq_len < 5 else [2, 3, 4, 5]
self.bleu = BLEU(test_text=tensor_to_tokens(self.gen_data.target,
self.index_word_dict),
real_text=tensor_to_tokens(
self.test_data.target,
self.test_data.index_word_dict),
gram=bleu_gram)
self.self_bleu = BLEU(
test_text=tensor_to_tokens(self.gen_data.target,
self.index_word_dict),
real_text=tensor_to_tokens(self.gen_data.target,
self.index_word_dict),
gram=3)
def __init__(self, opt):
self.log = create_logger(__name__, silent=False, to_disk=True,
log_file=cfg.log_filename if cfg.if_test
else [cfg.log_filename, cfg.save_root + 'log.txt'])
self.sig = Signal(cfg.signal_file)
self.opt = opt
self.show_config()
self.clas = None
# load dictionary
self.word2idx_dict, self.idx2word_dict = load_dict(cfg.dataset)
# Dataloader
try:
self.train_data = GenDataIter(cfg.train_data)
self.test_data = GenDataIter(cfg.test_data, if_test_data=True)
except:
pass
try:
self.train_data_list = [GenDataIter(cfg.cat_train_data.format(i)) for i in range(cfg.k_label)]
self.test_data_list = [GenDataIter(cfg.cat_test_data.format(i), if_test_data=True) for i in
range(cfg.k_label)]
self.clas_data_list = [GenDataIter(cfg.cat_test_data.format(str(i)), if_test_data=True) for i in
range(cfg.k_label)]
self.train_samples_list = [self.train_data_list[i].target for i in range(cfg.k_label)]
self.clas_samples_list = [self.clas_data_list[i].target for i in range(cfg.k_label)]
except:
pass
# Criterion
self.mle_criterion = nn.NLLLoss()
self.dis_criterion = nn.CrossEntropyLoss()
self.clas_criterion = nn.CrossEntropyLoss()
# Optimizer
self.clas_opt = None
# Metrics
self.bleu = BLEU('BLEU', gram=[2, 3, 4, 5], if_use=cfg.use_bleu)
self.nll_gen = NLL('NLL_gen', if_use=cfg.use_nll_gen, gpu=cfg.CUDA)
self.nll_div = NLL('NLL_div', if_use=cfg.use_nll_div, gpu=cfg.CUDA)
self.self_bleu = BLEU('Self-BLEU', gram=[2, 3, 4], if_use=cfg.use_self_bleu)
self.clas_acc = ACC(if_use=cfg.use_clas_acc)
self.ppl = PPL(self.train_data, self.test_data, n_gram=5, if_use=cfg.use_ppl)
self.all_metrics = [self.bleu, self.nll_gen, self.nll_div, self.self_bleu, self.ppl]
def __init__(self, opt):
super(RelGANInstructor, self).__init__(opt)
# generator, discriminator
self.gen = RelGAN_G(cfg.mem_slots,
cfg.num_heads,
cfg.head_size,
cfg.gen_embed_dim,
cfg.gen_hidden_dim,
cfg.vocab_size,
cfg.max_seq_len,
cfg.padding_idx,
gpu=cfg.CUDA)
self.dis = RelGAN_D(cfg.dis_embed_dim,
cfg.max_seq_len,
cfg.num_rep,
cfg.vocab_size,
cfg.padding_idx,
gpu=cfg.CUDA)
self.init_model()
# Optimizer
self.gen_opt = optim.Adam(self.gen.parameters(), lr=cfg.gen_lr)
self.gen_adv_opt = optim.Adam(self.gen.parameters(), lr=cfg.gen_adv_lr)
self.dis_opt = optim.Adam(self.dis.parameters(), lr=cfg.dis_lr)
# Criterion
self.mle_criterion = nn.NLLLoss()
# DataLoader
self.gen_data = GenDataIter(
self.gen.sample(cfg.batch_size, cfg.batch_size))
def getReward(self, samples):
"""
Get word-level reward and sentence-level reward of samples.
"""
batch_size, _ = samples.size()
inp, target = GenDataIter.prepare(
samples,
cfg.CUDA) # [batch_size, max_seq_len], [batch_size, max_seq_len]
inp = inp.transpose(1, 0) # [max_seq_len, batch_size]
target = target.transpose(1, 0) # [max_seq_len, batch_size]
dummy_tgt = torch.ones(self.max_seq_len, batch_size, dtype=torch.int)
if self.gpu:
dummy_tgt = dummy_tgt.cuda()
pred = self.forward(target, inp)
word_reward = F.nll_loss(pred,
target.contiguous().view(-1),
reduction='none').view(batch_size, -1)
sentence_reward = torch.mean(word_reward, dim=-1, keepdim=True)
#print(f"word reward len: {word_reward.size()} {word_reward}")
#print(f"word reward: {word_reward}")
#print(f"sentence reward len: {sentence_reward.size()}")
#print(f"sentence reward:{sentence_reward}")
return word_reward, sentence_reward
def __init__(self, opt):
super(RelGANInstructor, self).__init__(opt)
# generator, discriminator
self.gen = RelGAN_G(cfg.mem_slots, cfg.num_heads, cfg.head_size, cfg.gen_embed_dim, cfg.gen_hidden_dim,
cfg.vocab_size, cfg.max_seq_len, cfg.padding_idx, gpu=cfg.CUDA)
self.dis = RelGAN_D(cfg.dis_embed_dim, cfg.max_seq_len, cfg.num_rep, cfg.vocab_size, cfg.padding_idx,
gpu=cfg.CUDA)
self.init_model()
# Optimizer
self.gen_opt = optim.Adam(self.gen.parameters(), lr=cfg.gen_lr)
self.gen_adv_opt = optim.Adam(self.gen.parameters(), lr=cfg.gen_adv_lr)
self.dis_opt = optim.Adam(self.dis.parameters(), lr=cfg.dis_lr)
# Criterion
self.mle_criterion = nn.NLLLoss()
self.adv_criterion = nn.BCEWithLogitsLoss()
# DataLoader
self.gen_data = GenDataIter(self.gen.sample(cfg.batch_size, cfg.batch_size))
# Metrics
self.bleu = BLEU(test_text=tensor_to_tokens(self.gen_data.target, self.index_word_dict),
real_text=tensor_to_tokens(self.test_data.target, self.test_data.index_word_dict),
gram=[2, 3, 4, 5])
self.self_bleu = BLEU(test_text=tensor_to_tokens(self.gen_data.target, self.index_word_dict),
real_text=tensor_to_tokens(self.gen_data.target, self.index_word_dict),
gram=3)
def adv_train_generator(self):
g_loss = []
gen_data = GenDataIter(
self.old_gen.sample(cfg.samples_num, cfg.batch_size))
for (real, fake) in zip(self.oracle_data.loader, gen_data.loader):
real_inp, real_tar = real['input'], real['target']
fake_inp, fake_tar = fake['input'], fake['target']
if cfg.CUDA:
real_inp, real_tar, fake_inp, fake_tar = real_inp.cuda(
), real_tar.cuda(), fake_inp.cuda(), fake_tar.cuda()
# ===Train===
real_new_pred = self.cal_pred(self.gen, real_inp, real_tar)
real_old_pred = self.cal_pred(self.old_gen, real_inp, real_tar)
fake_new_pred = self.cal_pred(self.gen, fake_inp, fake_tar)
fake_old_pred = self.cal_pred(self.old_gen, fake_inp, fake_tar)
eps = 0
real_loss = -torch.sum(
torch.log(1 / (1 + real_old_pred /
(real_new_pred + eps) + eps) + eps))
fake_loss = -torch.sum(
torch.log(1 / (1 + fake_new_pred /
(fake_old_pred + eps) + eps) + eps))
adv_loss = real_loss + fake_loss
self.optimize(self.gen_adv_opt, adv_loss)
g_loss.append(adv_loss.item())
return np.mean(g_loss)
def cal_metrics(self, fmt_str=False):
"""
Calculate metrics
:param fmt_str: if return format string for logging
"""
with torch.no_grad():
# Prepare data for evaluation
eval_samples = self.gen.sample(cfg.samples_num, 4 * cfg.batch_size)
gen_data = GenDataIter(eval_samples)
gen_tokens = tensor_to_tokens(eval_samples, self.idx2word_dict)
#print(gen_tokens)
gen_tokens_s = tensor_to_tokens(self.gen.sample(200, 200),
self.idx2word_dict)
#print(gen_tokens_s)
# Reset metrics
self.bleu.reset(test_text=gen_tokens,
real_text=self.test_data.tokens)
self.nll_gen.reset(self.gen, self.train_data.loader)
self.nll_div.reset(self.gen, gen_data.loader)
self.self_bleu.reset(test_text=gen_tokens_s, real_text=gen_tokens)
self.ppl.reset(gen_tokens)
if fmt_str:
return ', '.join([
'%s = %s' % (metric.get_name(), metric.get_score())
for metric in self.all_metrics
])
else:
return [metric.get_score() for metric in self.all_metrics]
def __init__(self, opt):
self.log = create_logger(__name__,
silent=False,
to_disk=False if cfg.if_test else True,
log_file=None if cfg.if_test else
[cfg.log_filename, cfg.save_root + 'log.txt'])
self.sig = Signal(cfg.signal_file)
self.opt = opt
self.show_config()
# load dictionary
self.word_index_dict, self.index_word_dict = load_dict(cfg.dataset)
# Dataloader
self.oracle_data = GenDataIter(cfg.train_data)
self.test_data = GenDataIter(cfg.test_data)
def create_oracle():
"""Create a new Oracle model and Oracle's samples"""
import config as cfg
from models.Oracle import Oracle
print('Creating Oracle...')
oracle = Oracle(cfg.gen_embed_dim,
cfg.gen_hidden_dim,
cfg.vocab_size,
cfg.max_seq_len,
cfg.padding_idx,
gpu=cfg.CUDA)
if cfg.CUDA:
oracle = oracle.cuda()
torch.save(oracle.state_dict(), cfg.oracle_state_dict_path)
big_samples = oracle.sample(cfg.samples_num, 4 * cfg.batch_size)
# large
torch.save(big_samples, cfg.oracle_samples_path.format(cfg.samples_num))
# small
torch.save(oracle.sample(cfg.samples_num // 2, 4 * cfg.batch_size),
cfg.oracle_samples_path.format(cfg.samples_num // 2))
oracle_data = GenDataIter(big_samples)
mle_criterion = nn.NLLLoss()
groud_truth = NLL.cal_nll(oracle, oracle_data.loader, mle_criterion)
print('NLL_Oracle Groud Truth: %.4f' % groud_truth)
def cal_metrics(self, fmt_str=False):
with torch.no_grad():
# Prepare data for evaluation
eval_samples = self.gen.sample(cfg.samples_num, cfg.batch_size,
self.dis)
gen_data = GenDataIter(eval_samples)
gen_tokens = tensor_to_tokens(eval_samples, self.idx2word_dict)
gen_tokens_s = tensor_to_tokens(
self.gen.sample(200, cfg.batch_size, self.dis),
self.idx2word_dict)
# Reset metrics
self.bleu.reset(test_text=gen_tokens,
real_text=self.test_data.tokens)
self.nll_gen.reset(self.gen,
self.train_data.loader,
leak_dis=self.dis)
self.nll_div.reset(self.gen, gen_data.loader, leak_dis=self.dis)
self.self_bleu.reset(test_text=gen_tokens_s, real_text=gen_tokens)
self.ppl.reset(gen_tokens)
if fmt_str:
return ', '.join([
'%s = %s' % (metric.get_name(), metric.get_score())
for metric in self.all_metrics
])
else:
return [metric.get_score() for metric in self.all_metrics]
def adv_train_generator(self, g_step, current_k=0):
"""
The gen is trained using policy gradients, using the reward from the discriminator.
Training is done for num_batches batches.
"""
rollout_func = rollout.ROLLOUT(self.gen, cfg.CUDA)
adv_mana_loss = 0
adv_work_loss = 0
for step in range(g_step):
with torch.no_grad():
gen_samples = self.gen.sample(
cfg.batch_size, cfg.batch_size, self.dis,
train=True) # !!! train=True, the only place
inp, target = GenDataIter.prepare(gen_samples, gpu=cfg.CUDA)
# ===Train===
rewards = rollout_func.get_reward_leakgan(
target, cfg.rollout_num, self.dis,
current_k).cpu() # reward with MC search
mana_loss, work_loss = self.gen.adversarial_loss(
target, rewards, self.dis)
# update parameters
self.optimize_multi(self.gen_opt, [mana_loss, work_loss])
adv_mana_loss += mana_loss.data.item()
adv_work_loss += work_loss.data.item()
# ===Test===
self.log.info(
'[ADV-GEN] adv_mana_loss = %.4f, adv_work_loss = %.4f, %s' %
(adv_mana_loss / g_step, adv_work_loss / g_step,
self.cal_metrics(fmt_str=True)))
def adv_train_generator(self, g_step):
"""
The gen is trained using policy gradients, using the reward from the discriminator.
Training is done for num_batches batches.
"""
rollout_func = rollout.ROLLOUT(self.gen, cfg.CUDA)
total_g_loss = 0
print('g_step in adv_train_generator->', g_step)
for step in range(g_step):
samples = self.gen.sample(cfg.batch_size, cfg.batch_size)
print('samples ->', samples.size())
inp, target = GenDataIter.prepare(samples, gpu=cfg.CUDA)
print('inp ->', inp.size())
print('target ->', target.size())
# ===Train===
rewards = rollout_func.get_reward(target, cfg.rollout_num, self.dis)
print('rewards ->', rewards.size(), rewards)
adv_loss = self.gen.batchPGLoss(inp, target, rewards)
self.optimize(self.gen_adv_opt, adv_loss)
total_g_loss += adv_loss.item()
# ===Test===
self.log.info('[ADV-GEN]: g_loss = %.4f, %s' % (total_g_loss, self.cal_metrics(fmt_str=True)))
def create_multi_oracle(number):
for i in range(number):
print('Creating Oracle %d...' % i)
oracle = Oracle(cfg.gen_embed_dim,
cfg.gen_hidden_dim,
cfg.vocab_size,
cfg.max_seq_len,
cfg.padding_idx,
gpu=cfg.CUDA)
if cfg.CUDA:
oracle = oracle.cuda()
large_samples = oracle.sample(cfg.samples_num, 4 * cfg.batch_size)
small_samples = oracle.sample(cfg.samples_num // 2, 4 * cfg.batch_size)
torch.save(oracle.state_dict(),
cfg.multi_oracle_state_dict_path.format(i))
torch.save(large_samples,
cfg.multi_oracle_samples_path.format(i, cfg.samples_num))
torch.save(
small_samples,
cfg.multi_oracle_samples_path.format(i, cfg.samples_num // 2))
oracle_data = GenDataIter(large_samples)
mle_criterion = nn.NLLLoss()
groud_truth = NLL.cal_nll(oracle, oracle_data.loader, mle_criterion)
print('Oracle %d Groud Truth: %.4f' % (i, groud_truth))
def create_many_oracle(from_a, to_b, num=1, save_path='../pretrain/'):
for i in range(num):
while True:
oracle = Oracle(cfg.gen_embed_dim,
cfg.gen_hidden_dim,
cfg.vocab_size,
cfg.max_seq_len,
cfg.padding_idx,
gpu=cfg.CUDA)
if cfg.CUDA:
oracle = oracle.cuda()
big_samples = oracle.sample(cfg.samples_num, 8 * cfg.batch_size)
small_samples = oracle.sample(cfg.samples_num // 2,
8 * cfg.batch_size)
oracle_data = GenDataIter(big_samples)
mle_criterion = nn.NLLLoss()
groud_truth = NLL.cal_nll(oracle, oracle_data.loader,
mle_criterion)
if from_a <= groud_truth <= to_b:
print('save ground truth: ', groud_truth)
prefix = 'oracle_lstm'
torch.save(oracle.state_dict(),
save_path + '{}.pt'.format(prefix))
torch.save(
big_samples, save_path +
'{}_samples_{}.pt'.format(prefix, cfg.samples_num))
torch.save(
small_samples, save_path +
'{}_samples_{}.pt'.format(prefix, cfg.samples_num // 2))
break
def __init__(self, opt):
self.log = create_logger(__name__, silent=False, to_disk=True,
log_file=cfg.log_filename if cfg.if_test
else [cfg.log_filename, cfg.save_root + 'log.txt'])
self.sig = Signal(cfg.signal_file)
self.opt = opt
# oracle, generator, discriminator
self.oracle = Oracle(cfg.gen_embed_dim, cfg.gen_hidden_dim, cfg.vocab_size, cfg.max_seq_len,
cfg.padding_idx, gpu=cfg.CUDA)
self.dis = None
self.show_config()
# DataLoader
if not os.path.exists(cfg.oracle_samples_path.format(cfg.samples_num)) or not cfg.oracle_pretrain:
create_oracle()
self.oracle.load_state_dict(torch.load(cfg.oracle_state_dict_path))
self.oracle_samples = torch.load(cfg.oracle_samples_path.format(cfg.samples_num))
self.oracle_data = GenDataIter(self.oracle_samples)
self.gen_data = None
self.dis_data = None
# Criterion
self.mle_criterion = nn.NLLLoss()
self.dis_criterion = None
def cal_metrics_with_label(self, label_i):
assert type(label_i) == int, 'missing label'
with torch.no_grad():
# Prepare data for evaluation
eval_samples = self.gen.sample(cfg.samples_num,
8 * cfg.batch_size,
label_i=label_i)
gen_data = GenDataIter(eval_samples)
gen_tokens = tensor_to_tokens(eval_samples, self.idx2word_dict)
gen_tokens_s = tensor_to_tokens(
self.gen.sample(200, 200, label_i=label_i), self.idx2word_dict)
clas_data = CatClasDataIter([eval_samples], label_i)
# Reset metrics
self.bleu.reset(test_text=gen_tokens,
real_text=self.test_data_list[label_i].tokens)
self.nll_gen.reset(self.gen, self.train_data_list[label_i].loader,
label_i)
self.nll_div.reset(self.gen, gen_data.loader, label_i)
self.self_bleu.reset(test_text=gen_tokens_s, real_text=gen_tokens)
self.clas_acc.reset(self.clas, clas_data.loader)
self.ppl.reset(gen_tokens)
return [metric.get_score() for metric in self.all_metrics]
def __init__(self, opt):
super(LeakGANInstructor, self).__init__(opt)
# generator, discriminator
self.gen = LeakGAN_G(cfg.gen_embed_dim, cfg.gen_hidden_dim,
cfg.vocab_size, cfg.max_seq_len, cfg.padding_idx,
cfg.goal_size, cfg.step_size, cfg.CUDA)
self.dis = LeakGAN_D(cfg.dis_embed_dim,
cfg.vocab_size,
cfg.padding_idx,
gpu=cfg.CUDA)
#LSTM
self.corpus = dataa.Corpus('dataset/emnlp_news/')
self.lstm = LSTM.RNNModel('LSTM', len(self.corpus.dictionary), 200,
600, 3, 0.2, False)
if (cfg.CUDA):
self.dis.cuda()
self.gen.cuda()
self.init_model()
# optimizer
mana_params, work_params = self.gen.split_params()
mana_opt = optim.Adam(mana_params, lr=cfg.gen_lr)
work_opt = optim.Adam(work_params, lr=cfg.gen_lr)
self.gen_opt = [mana_opt, work_opt]
self.dis_opt = optim.Adam(self.dis.parameters(), lr=cfg.dis_lr)
# Criterion
self.mle_criterion = nn.NLLLoss()
self.dis_criterion = nn.CrossEntropyLoss()
# DataLoader
self.gen_data = GenDataIter(
self.gen.sample(cfg.batch_size, cfg.batch_size, self.dis))
self.dis_data = DisDataIter(self.gen_data.random_batch()['target'],
self.oracle_data.random_batch()['target'])
# Metrics
self.bleu3 = BLEU(test_text=tensor_to_tokens(self.gen_data.target,
self.index_word_dict),
real_text=tensor_to_tokens(self.test_data.target,
self.index_word_dict),
gram=3)
def cal_metrics_with_label(self, label_i):
assert type(label_i) == int, 'missing label'
with torch.no_grad():
# Prepare data for evaluation
eval_samples = self.gen.sample(cfg.samples_num, 8 * cfg.batch_size, label_i=label_i)
gen_data = GenDataIter(eval_samples)
# Reset metrics
self.nll_oracle.reset(self.oracle_list[label_i], gen_data.loader, label_i)
self.nll_gen.reset(self.gen, self.oracle_data_list[label_i].loader, label_i)
self.nll_div.reset(self.gen, gen_data.loader, label_i)
return [metric.get_score() for metric in self.all_metrics]
def __init__(self, opt):
self.log = create_logger(__name__,
silent=False,
to_disk=True,
log_file=cfg.log_filename if cfg.if_test else
[cfg.log_filename, cfg.save_root + 'log.txt'])
self.sig = Signal(cfg.signal_file)
self.opt = opt
self.show_config()
# load dictionary
self.word_index_dict, self.index_word_dict = load_dict(cfg.dataset)
# Dataloader
self.train_data = GenDataIter(cfg.train_data)
self.test_data = GenDataIter(cfg.test_data, if_test_data=True)
self.gen_data = None
# Criterion
self.mle_criterion = nn.NLLLoss()
self.dis_criterion = None
self.bleu = None
self.self_bleu = None
def getReward(self, samples):
"""
Get word-level reward and sentence-level reward of samples.
"""
batch_size, _ = samples.size()
inp, target = GenDataIter.prepare(samples, cfg.CUDA)
hidden = self.init_hidden(batch_size)
pred = self.forward(inp, hidden)
word_reward = F.nll_loss(pred, target.view(-1),
reduction='none').view(batch_size, -1)
sentence_reward = torch.mean(word_reward, dim=-1, keepdim=True)
return word_reward, sentence_reward
def getReward(self, samples):
"""
Get word-level reward and sentence-level reward of samples.
"""
batch_size, _ = samples.size()
inp, target = GenDataIter.prepare(samples, cfg.CUDA)
src_mask = self.generate_square_subsequent_mask(self.max_seq_len)
#pred = self.forward(inp, src_mask)
pred = self.forward(inp)
word_reward = F.nll_loss(pred, target.view(-1),
reduction='none').view(batch_size, -1)
sentence_reward = torch.mean(word_reward, dim=-1, keepdim=True)
return word_reward, sentence_reward
def adv_train_generator(self, g_step):
"""
Train the generator with mediator rewards
"""
g_loss = []
for step in range(g_step):
inp, target = GenDataIter.prepare(self.gen.sample(
cfg.batch_size, cfg.batch_size),
gpu=cfg.CUDA)
# ===Train===
rewards = self.dis(inp, self.dis.init_hidden(cfg.batch_size))
loss = self.gen.get_loss(inp, rewards)
self.optimize(self.gen_adv_opt, loss)
g_loss.append(loss.item())
return np.mean(g_loss)
def cal_metrics(self, fmt_str=False):
"""
Calculate metrics
:param fmt_str: if return format string for logging
"""
with torch.no_grad():
# Prepare data for evaluation
gen_data = GenDataIter(self.gen.sample(cfg.samples_num, 4 * cfg.batch_size))
# Reset metrics
self.nll_oracle.reset(self.oracle, gen_data.loader)
self.nll_gen.reset(self.gen, self.oracle_data.loader)
self.nll_div.reset(self.gen, gen_data.loader)
if fmt_str:
return ', '.join(['%s = %s' % (metric.get_name(), metric.get_score()) for metric in self.all_metrics])
else:
return [metric.get_score() for metric in self.all_metrics]
def cal_metrics(self, fmt_str=False):
# Prepare data for evaluation
gen_data = GenDataIter(
self.gen.sample(cfg.samples_num, cfg.batch_size, self.dis))
# Reset metrics
self.nll_oracle.reset(self.oracle, gen_data.loader)
self.nll_gen.reset(self.gen,
self.oracle_data.loader,
leak_dis=self.dis)
self.nll_div.reset(self.gen, gen_data.loader, leak_dis=self.dis)
if fmt_str:
return ', '.join([
'%s = %s' % (metric.get_name(), metric.get_score())
for metric in self.all_metrics
])
else:
return [metric.get_score() for metric in self.all_metrics]
def adv_train_generator(self, g_step):
"""
The gen is trained by MLE-like objective.
"""
total_g_loss = 0
for step in range(g_step):
inp, target = GenDataIter.prepare(self.gen.sample(
cfg.batch_size, cfg.batch_size),
gpu=cfg.CUDA)
# ===Train===
rewards = self.get_mali_reward(target)
adv_loss = self.gen.adv_loss(inp, target, rewards)
self.optimize(self.gen_adv_opt, adv_loss)
total_g_loss += adv_loss.item()
# ===Test===
self.log.info('[ADV-GEN]: g_loss = %.4f, %s' %
(total_g_loss, self.cal_metrics(fmt_str=True)))
