def train_GAN(conf_data):
"""Training Process for GAN.
Parameters
----------
conf_data: dict
Dictionary containing all parameters and objects.
Returns
-------
conf_data: dict
Dictionary containing all parameters and objects.
"""
seq = conf_data['GAN_model']['seq']
if seq == 1:
pre_epoch_num = conf_data['generator']['pre_epoch_num']
GENERATED_NUM = 10000
EVAL_FILE = 'eval.data'
POSITIVE_FILE = 'real.data'
NEGATIVE_FILE = 'gene.data'
temp = 1 #TODO Determines how many times is the discriminator updated. Take this as a value input
epochs = int(conf_data['GAN_model']['epochs'])
if seq == 0:
dataloader = conf_data['data_learn']
mini_batch_size = int(conf_data['GAN_model']['mini_batch_size'])
data_label = int(conf_data['GAN_model']['data_label'])
cuda = conf_data['cuda']
g_latent_dim = int(conf_data['generator']['latent_dim'])
classes = int(conf_data['GAN_model']['classes'])
w_loss = int(conf_data['GAN_model']['w_loss'])
clip_value = float(conf_data['GAN_model']['clip_value'])
n_critic = int(conf_data['GAN_model']['n_critic'])
lambda_gp = int(conf_data['GAN_model']['lambda_gp'])
log_file = open(conf_data['performance_log'] + "/log.txt", "w+")
#Covert these to parameters of the config data
Tensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor
conf_data['Tensor'] = Tensor
LongTensor = torch.cuda.LongTensor if cuda else torch.LongTensor
conf_data['LongTensor'] = LongTensor
FloatTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor
conf_data['FloatTensor'] = FloatTensor
conf_data['epochs'] = epochs
#print ("Just before training")
if seq == 1: #TODO: Change back to 1
target_lstm = TargetLSTM(conf_data['GAN_model']['vocab_size'],
conf_data['generator']['embedding_dim'],
conf_data['generator']['hidden_dim'],
conf_data['cuda'])
if cuda == True:
target_lstm = target_lstm.cuda()
conf_data['target_lstm'] = target_lstm
gen_data_iter = GenDataIter('real.data', mini_batch_size)
generator = conf_data['generator_model']
discriminator = conf_data['discriminator_model']
g_loss_func = conf_data['generator_loss']
d_loss_func = conf_data['discriminator_loss']
optimizer_D = conf_data['discriminator_optimizer']
optimizer_G = conf_data['generator_optimizer']
#print('Pretrain with MLE ...')
for epoch in range(pre_epoch_num): #TODO: Change the range
loss = train_epoch(generator, gen_data_iter, g_loss_func,
optimizer_G, conf_data, 'g')
print('Epoch [%d] Model Loss: %f' % (epoch, loss))
generate_samples(generator, mini_batch_size, GENERATED_NUM,
EVAL_FILE, conf_data)
eval_iter = GenDataIter(EVAL_FILE, mini_batch_size)
loss = eval_epoch(target_lstm, eval_iter, g_loss_func, conf_data)
print('Epoch [%d] True Loss: %f' % (epoch, loss))
dis_criterion = d_loss_func
dis_optimizer = optimizer_D
#TODO: Understand why the below two code line were there ?
# if conf_data['cuda']:
# dis_criterion = dis_criterion.cuda()
#print('Pretrain Dsicriminator ...')
dis_data_iter = DisDataIter(POSITIVE_FILE, NEGATIVE_FILE,
mini_batch_size)
for epoch in range(5): #TODO: change back 5
generate_samples(generator, mini_batch_size, GENERATED_NUM,
NEGATIVE_FILE, conf_data)
dis_data_iter = DisDataIter(POSITIVE_FILE, NEGATIVE_FILE,
mini_batch_size)
for _ in range(3): #TODO: change back 3
loss = train_epoch(discriminator, dis_data_iter, dis_criterion,
dis_optimizer, conf_data, 'd')
print('Epoch [%d], loss: %f' % (epoch, loss))
conf_data['generator_model'] = generator
conf_data['discriminator_model'] = discriminator
torch.save(conf_data['generator_model'].state_dict(),
conf_data['save_model_path'] + '/Seq/' + 'pre_generator.pt')
torch.save(
conf_data['discriminator_model'].state_dict(),
conf_data['save_model_path'] + '/Seq/' + 'pre_discriminator.pt')
conf_data['rollout'] = Rollout(generator, 0.8)
for epoch in range(epochs):
conf_data['epoch'] = epoch
if seq == 0:
to_iter = dataloader
elif seq == 1: #TODO: Change this back to 1
to_iter = [1]
for i, iterator in enumerate(to_iter):
optimizer_D = conf_data['discriminator_optimizer']
optimizer_G = conf_data['generator_optimizer']
generator = conf_data['generator_model']
discriminator = conf_data['discriminator_model']
g_loss_func = conf_data['generator_loss']
d_loss_func = conf_data['discriminator_loss']
# if aux = 1:
#print ("Reached here --------------> ")
conf_data['iterator'] = i
if seq == 0:
if data_label == 1:
imgs, labels = iterator
else:
imgs = iterator
# Adversarial ground truths
valid = Variable(Tensor(imgs.size(0), 1).fill_(1.0),
requires_grad=False)
conf_data['valid'] = valid
fake = Variable(Tensor(imgs.size(0), 1).fill_(0.0),
requires_grad=False)
conf_data['fake'] = fake
# Configure input
real_imgs = Variable(imgs.type(Tensor))
if data_label == 1:
labels = Variable(labels.type(LongTensor))
# Sample noise as generator input
z = Variable(
Tensor(
np.random.normal(0, 1, (imgs.shape[0], g_latent_dim))))
if classes > 0:
gen_labels = Variable(
LongTensor(np.random.randint(0, classes,
imgs.shape[0])))
conf_data['gen_labels'] = gen_labels
# elif seq == 1: #If yes seqGAN
# # samples = generator.sample(mini_batch_size,conf_data['generator']['sequece_length'])
# # zeros = torch.zeros((mini_batch_size,1)).type(LongTensor)
# # imgs = Variable(torch.cat([zeros,samples.data]),dim=1)[:,:-1].contiguous() #TODO: change imgs to inps all, to make more sense of the code
# # targets = Variable(sample.data).contiguous().view((-1,))
# # rewards = rollout.get_reward(sample,16,discriminator)
# # rewards = Variable(Tensor(rewards))
# # prob = generator.forward(inputs)
# # loss = gen_gan_loss(prob)
# pass
# #optimizer_G
# ---------------------
# Train Discriminator
# ---------------------
optimizer_D.zero_grad()
if seq == 1: #TODO change this back to 1
dis_data_iter = DisDataIter(POSITIVE_FILE, NEGATIVE_FILE,
mini_batch_size)
for i in range(
temp
): # TODO: Make this a parameter -> for x updates --> I am read the stored models here as well. Should I reamove this ???
optimizer_D = conf_data['discriminator_optimizer']
optimizer_G = conf_data['generator_optimizer']
generator = conf_data['generator_model']
discriminator = conf_data['discriminator_model']
g_loss_func = conf_data['generator_loss']
d_loss_func = conf_data['discriminator_loss']
if classes <= 0:
#print ("Reached here 2 --------------> ")
if seq == 0:
gen_imgs = generator(z)
# Measure discriminator's ability to classify real from generated samples
#Real images
real_validity = discriminator(real_imgs)
#Fake images
fake_validity = discriminator(gen_imgs.detach())
if seq == 1:
generate_samples(generator, mini_batch_size,
GENERATED_NUM, NEGATIVE_FILE,
conf_data)
dis_data_iter = DisDataIter(POSITIVE_FILE,
NEGATIVE_FILE,
mini_batch_size)
loss = train_epoch(discriminator, dis_data_iter,
d_loss_func, optimizer_D, conf_data,
'd')
conf_data['d_loss'] = loss
#exit()
else:
if seq == 0:
gen_imgs = generator(z, gen_labels)
real_validity = discriminator(real_imgs, labels)
fake_validity = discriminator(gen_imgs.detach(),
labels)
if seq == 0:
conf_data['gen_imgs'] = gen_imgs
if seq == 0:
if w_loss == 0:
real_loss = d_loss_func.loss(real_validity, valid)
fake_loss = d_loss_func.loss(fake_validity, fake)
d_loss = (real_loss + fake_loss) / 2
elif w_loss == 1:
d_loss = -d_loss_func.loss(real_validity,
valid) + d_loss_func.loss(
fake_validity, fake)
if lambda_gp > 0:
conf_data['real_data_sample'] = real_imgs.data
conf_data['fake_data_sample'] = gen_imgs.data
conf_data = compute_gradient_penalty(conf_data)
gradient_penalty = conf_data['gradient_penalty']
d_loss = d_loss + lambda_gp * gradient_penalty
conf_data['d_loss'] = d_loss
d_loss.backward()
optimizer_D.step()
if clip_value > 0:
# Clip weights of discriminator
for p in discriminator.parameters():
p.data.clamp_(-clip_value, clip_value)
# -----------------
# Train Generator
# -----------------
conf_data['generator_model'] = generator
conf_data['discriminator_model'] = discriminator
#Next 4 lines were recently added maybe have to remove this.
conf_data['optimizer_G'] = optimizer_G
conf_data['optimizer_D'] = optimizer_D
conf_data['generator_loss'] = g_loss_func
conf_data['discriminator_loss'] = d_loss_func
if seq == 0:
conf_data['noise'] = z
if n_critic <= 0:
conf_data = training_fucntion_generator(conf_data)
elif n_critic > 0:
# Train the generator every n_critic iterations
if i % n_critic == 0:
conf_data = training_fucntion_generator(conf_data)
#exit()
# print ("------------------ Here (train_GAN.py)")
if seq == 0:
batches_done = epoch * len(dataloader) + i
if batches_done % int(conf_data['sample_interval']) == 0:
if classes <= 0:
# print ("Here")
# print (type(gen_imgs.data[:25]))
# print (gen_imgs.data[:25].shape)
save_image(gen_imgs.data[:25],
conf_data['result_path'] +
'/%d.png' % batches_done,
nrow=5,
normalize=True)
elif classes > 0:
sample_image(10, batches_done, conf_data)
if seq == 0:
log_file.write("[Epoch %d/%d] [D loss: %f] [G loss: %f] \n" %
(epoch, epochs, conf_data['d_loss'].item(),
conf_data['g_loss'].item()))
elif seq == 1:
# print ("Done")
log_file.write(
"[Epoch %d/%d] [D loss: %f] [G loss: %f] \n" %
(epoch, epochs, conf_data['d_loss'], conf_data['g_loss']))
conf_data['generator_model'] = generator
conf_data['discriminator_model'] = discriminator
conf_data['log_file'] = log_file
return conf_data
def main():
random.seed(SEED)
np.random.seed(SEED)
# Define Networks
generator = Generator(VOCAB_SIZE, g_emb_dim, g_hidden_dim, opt.cuda)
discriminator = Discriminator(d_num_class, VOCAB_SIZE, d_emb_dim,
d_filter_sizes, d_num_filters, d_dropout)
target_lstm = TargetLSTM(VOCAB_SIZE, g_emb_dim, g_hidden_dim, opt.cuda)
if opt.cuda:
generator = generator.cuda()
discriminator = discriminator.cuda()
target_lstm = target_lstm.cuda()
# Generate toy data using target lstm
print('Generating data ...')
generate_samples(target_lstm, BATCH_SIZE, GENERATED_NUM, POSITIVE_FILE)
# Load data from file
gen_data_iter = GenDataIter(POSITIVE_FILE, BATCH_SIZE)
# Pretrain Generator using MLE
gen_criterion = nn.NLLLoss(size_average=False)
gen_optimizer = optim.Adam(generator.parameters())
if opt.cuda:
gen_criterion = gen_criterion.cuda()
print('Pretrain with MLE ...')
for epoch in range(PRE_EPOCH_NUM):
loss = train_epoch(generator, gen_data_iter, gen_criterion,
gen_optimizer)
print('Epoch [%d] Model Loss: %f' % (epoch, loss))
generate_samples(generator, BATCH_SIZE, GENERATED_NUM, EVAL_FILE)
eval_iter = GenDataIter(EVAL_FILE, BATCH_SIZE)
loss = eval_epoch(target_lstm, eval_iter, gen_criterion)
print('Epoch [%d] True Loss: %f' % (epoch, loss))
# Pretrain Discriminator
dis_criterion = nn.NLLLoss(size_average=False)
dis_optimizer = optim.Adam(discriminator.parameters())
if opt.cuda:
dis_criterion = dis_criterion.cuda()
print('Pretrain Dsicriminator ...')
for epoch in range(5):
generate_samples(generator, BATCH_SIZE, GENERATED_NUM, NEGATIVE_FILE)
dis_data_iter = DisDataIter(POSITIVE_FILE, NEGATIVE_FILE, BATCH_SIZE)
for _ in range(3):
loss = train_epoch(discriminator, dis_data_iter, dis_criterion,
dis_optimizer)
print('Epoch [%d], loss: %f' % (epoch, loss))
# Adversarial Training
rollout = Rollout(generator, 0.8)
print('#####################################################')
print('Start Adeversatial Training...\n')
gen_gan_loss = GANLoss()
gen_gan_optm = optim.Adam(generator.parameters())
if opt.cuda:
gen_gan_loss = gen_gan_loss.cuda()
gen_criterion = nn.NLLLoss(size_average=False)
if opt.cuda:
gen_criterion = gen_criterion.cuda()
dis_criterion = nn.NLLLoss(size_average=False)
dis_optimizer = optim.Adam(discriminator.parameters())
if opt.cuda:
dis_criterion = dis_criterion.cuda()
for total_batch in range(TOTAL_BATCH):
# Train the generator for one step
for it in range(1):
samples = generator.sample(BATCH_SIZE, g_sequence_len)
# construct the input to the generator, add zeros before samples and delete the last column
zeros = torch.zeros((BATCH_SIZE, 1)).type(torch.LongTensor)
if samples.is_cuda:
zeros = zeros.cuda()
inputs = Variable(
torch.cat([zeros, samples.data], dim=1)[:, :-1].contiguous())
targets = Variable(samples.data).contiguous().view((-1, ))
# calculate the reward
rewards = rollout.get_reward(samples, 16, discriminator)
rewards = Variable(torch.Tensor(rewards))
if opt.cuda:
rewards = torch.exp(rewards.cuda()).contiguous().view((-1, ))
prob = generator.forward(inputs)
loss = gen_gan_loss(prob, targets, rewards)
gen_gan_optm.zero_grad()
loss.backward()
gen_gan_optm.step()
if total_batch % 1 == 0 or total_batch == TOTAL_BATCH - 1:
generate_samples(generator, BATCH_SIZE, GENERATED_NUM, EVAL_FILE)
eval_iter = GenDataIter(EVAL_FILE, BATCH_SIZE)
loss = eval_epoch(target_lstm, eval_iter, gen_criterion)
print('Batch [%d] True Loss: %f' % (total_batch, loss))
rollout.update_params()
for _ in range(4):
generate_samples(generator, BATCH_SIZE, GENERATED_NUM,
NEGATIVE_FILE)
dis_data_iter = DisDataIter(POSITIVE_FILE, NEGATIVE_FILE,
BATCH_SIZE)
for _ in range(2):
loss = train_epoch(discriminator, dis_data_iter, dis_criterion,
dis_optimizer)
torch.cuda.manual_seed(args.seed)
if not args.hpc:
args.data_path = 'data_test/'
POSITIVE_FILE = args.data_path + POSITIVE_FILE
NEGATIVE_FILE = args.data_path + NEGATIVE_FILE
EPOCH_FILE = args.data_path + EPOCH_FILE
if os.path.exists(EPOCH_FILE):
os.remove(EPOCH_FILE)
# Set models, criteria, optimizers
generator = Generator(args.vocab_size, g_embed_dim, g_hidden_dim,
args.cuda)
discriminator = Discriminator(d_num_class, args.vocab_size, d_embed_dim,
d_filter_sizes, d_num_filters,
d_dropout_prob)
target_lstm = TargetLSTM(args.vocab_size, g_embed_dim, g_hidden_dim,
args.cuda)
nll_loss = nn.NLLLoss()
pg_loss = PGLoss()
if args.cuda:
generator = generator.cuda()
discriminator = discriminator.cuda()
target_lstm = target_lstm.cuda()
nll_loss = nll_loss.cuda()
pg_loss = pg_loss.cuda()
cudnn.benchmark = True
gen_optimizer = optim.Adam(params=generator.parameters(), lr=args.gen_lr)
dis_optimizer = optim.SGD(params=discriminator.parameters(),
lr=args.dis_lr)
# Container of experiment data
gen_pretrain_train_loss = []
def main():
random.seed(SEED)
np.random.seed(SEED)
calc_bleu([1, 10, 12])
exit()
# Build up dataset
s_train, s_test = load_from_big_file('../data/train_data_obama.txt')
# idx_to_word: List of id to word
# word_to_idx: Dictionary mapping word to id
idx_to_word, word_to_idx = fetch_vocab(s_train, s_train, s_test)
# TODO: 1. Prepare data for attention model
# input_seq, target_seq = prepare_data(DATA_GERMAN, DATA_ENGLISH, word_to_idx)
global VOCAB_SIZE
VOCAB_SIZE = len(idx_to_word)
save_vocab(CHECKPOINT_PATH + 'metadata.data', idx_to_word, word_to_idx,
VOCAB_SIZE, g_emb_dim, g_hidden_dim, g_sequence_len)
print('VOCAB SIZE:', VOCAB_SIZE)
# Define Networks
generator = Generator(VOCAB_SIZE, g_emb_dim, g_hidden_dim, g_sequence_len,
BATCH_SIZE, opt.cuda)
discriminator = Discriminator(d_num_class, VOCAB_SIZE, d_emb_dim,
d_filter_sizes, d_num_filters, d_dropout)
target_lstm = TargetLSTM(VOCAB_SIZE, g_emb_dim, g_hidden_dim, opt.cuda)
if opt.cuda:
generator = generator.cuda()
discriminator = discriminator.cuda()
target_lstm = target_lstm.cuda()
# Generate toy data using target lstm
print('Generating data ...')
generate_real_data('../data/train_data_obama.txt', BATCH_SIZE,
GENERATED_NUM, idx_to_word, word_to_idx, POSITIVE_FILE,
TEST_FILE)
# Create Test data iterator for testing
test_iter = GenDataIter(TEST_FILE, BATCH_SIZE)
# generate_samples(target_lstm, BATCH_SIZE, GENERATED_NUM, POSITIVE_FILE, idx_to_word)
# Load data from file
gen_data_iter = GenDataIter(POSITIVE_FILE, BATCH_SIZE)
# Pretrain Generator using MLE
# gen_criterion = nn.NLLLoss(size_average=False)
gen_criterion = nn.CrossEntropyLoss()
gen_optimizer = optim.Adam(generator.parameters())
if opt.cuda:
gen_criterion = gen_criterion.cuda()
print('Pretrain with MLE ...')
for epoch in range(PRE_EPOCH_NUM):
loss = train_epoch(generator, gen_data_iter, gen_criterion,
gen_optimizer)
print('Epoch [%d] Model Loss: %f' % (epoch, loss))
print('Training Output')
test_predict(generator, test_iter, idx_to_word, train_mode=True)
sys.stdout.flush()
# TODO: 2. Flags to ensure dimension of model input is handled
# generate_samples(generator, BATCH_SIZE, GENERATED_NUM, EVAL_FILE)
"""
eval_iter = GenDataIter(EVAL_FILE, BATCH_SIZE)
print('Iterator Done')
loss = eval_epoch(target_lstm, eval_iter, gen_criterion)
print('Epoch [%d] True Loss: %f' % (epoch, loss))
"""
print('OUTPUT AFTER PRE-TRAINING')
test_predict(generator, test_iter, idx_to_word, train_mode=True)
# Pretrain Discriminator
dis_criterion = nn.NLLLoss(size_average=False)
dis_optimizer = optim.Adam(discriminator.parameters())
if opt.cuda:
dis_criterion = dis_criterion.cuda()
print('Pretrain Discriminator ...')
for epoch in range(3):
generate_samples(generator, BATCH_SIZE, GENERATED_NUM, NEGATIVE_FILE)
dis_data_iter = DisDataIter(POSITIVE_FILE, NEGATIVE_FILE, BATCH_SIZE)
for _ in range(3):
loss = train_epoch(discriminator, dis_data_iter, dis_criterion,
dis_optimizer)
print('Epoch [%d], loss: %f' % (epoch, loss))
sys.stdout.flush()
# Adversarial Training
rollout = Rollout(generator, 0.8)
print('#####################################################')
print('Start Adversarial Training...\n')
gen_gan_loss = GANLoss()
gen_gan_optm = optim.Adam(generator.parameters())
if opt.cuda:
gen_gan_loss = gen_gan_loss.cuda()
gen_criterion = nn.NLLLoss(size_average=False)
if opt.cuda:
gen_criterion = gen_criterion.cuda()
dis_criterion = nn.NLLLoss(size_average=False)
dis_optimizer = optim.Adam(discriminator.parameters())
if opt.cuda:
dis_criterion = dis_criterion.cuda()
real_iter = GenDataIter(POSITIVE_FILE, BATCH_SIZE)
for total_batch in range(TOTAL_BATCH):
## Train the generator for one step
for it in range(1):
if real_iter.idx >= real_iter.data_num:
real_iter.reset()
inputs = real_iter.next()[0]
inputs = inputs.cuda()
samples = generator.sample(BATCH_SIZE, g_sequence_len, inputs)
samples = samples.cpu()
rewards = rollout.get_reward(samples, 16, discriminator)
rewards = Variable(torch.Tensor(rewards))
if opt.cuda:
rewards = torch.exp(rewards.cuda()).contiguous().view((-1, ))
prob = generator.forward(inputs)
mini_batch = prob.shape[0]
prob = torch.reshape(
prob,
(prob.shape[0] * prob.shape[1], -1)) #prob.view(-1, g_emb_dim)
targets = copy.deepcopy(inputs).contiguous().view((-1, ))
loss = gen_gan_loss(prob, targets, rewards)
gen_gan_optm.zero_grad()
loss.backward()
gen_gan_optm.step()
"""
samples = generator.sample(BATCH_SIZE, g_sequence_len)
# construct the input to the genrator, add zeros before samples and delete the last column
zeros = torch.zeros((BATCH_SIZE, 1)).type(torch.LongTensor)
if samples.is_cuda:
zeros = zeros.cuda()
inputs = Variable(torch.cat([zeros, samples.data], dim = 1)[:, :-1].contiguous())
targets = Variable(samples.data).contiguous().view((-1,))
print('', inputs.shape, targets.shape)
print(inputs, targets)
# calculate the reward
rewards = rollout.get_reward(samples, 16, discriminator)
rewards = Variable(torch.Tensor(rewards))
if opt.cuda:
rewards = torch.exp(rewards.cuda()).contiguous().view((-1,))
prob = generator.forward(inputs)
mini_batch = prob.shape[0]
prob = torch.reshape(prob, (prob.shape[0] * prob.shape[1], -1)) #prob.view(-1, g_emb_dim)
loss = gen_gan_loss(prob, targets, rewards)
gen_gan_optm.zero_grad()
loss.backward()
gen_gan_optm.step()
"""
print('Batch [%d] True Loss: %f' % (total_batch, loss))
if total_batch % 1 == 0 or total_batch == TOTAL_BATCH - 1:
# generate_samples(generator, BATCH_SIZE, GENERATED_NUM, EVAL_FILE)
# eval_iter = GenDataIter(EVAL_FILE, BATCH_SIZE)
# loss = eval_epoch(target_lstm, eval_iter, gen_criterion)
if len(prob.shape) > 2:
prob = torch.reshape(prob, (prob.shape[0] * prob.shape[1], -1))
predictions = torch.max(prob, dim=1)[1]
predictions = predictions.view(mini_batch, -1)
for each_sen in list(predictions):
print('Train Output:',
generate_sentence_from_id(idx_to_word, each_sen))
test_predict(generator, test_iter, idx_to_word, train_mode=True)
torch.save(generator.state_dict(),
CHECKPOINT_PATH + 'generator.model')
torch.save(discriminator.state_dict(),
CHECKPOINT_PATH + 'discriminator.model')
rollout.update_params()
for _ in range(4):
generate_samples(generator, BATCH_SIZE, GENERATED_NUM,
NEGATIVE_FILE)
dis_data_iter = DisDataIter(POSITIVE_FILE, NEGATIVE_FILE,
BATCH_SIZE)
for _ in range(2):
loss = train_epoch(discriminator, dis_data_iter, dis_criterion,
dis_optimizer)
C = prob.size(1)
one_hot = torch.zeros((N, C))
if prob.is_cuda:
one_hot = one_hot.cuda()
one_hot.scatter_(1, target.data.view((-1, 1)), 1)
one_hot = one_hot.type(torch.ByteTensor)
one_hot = Variable(one_hot)
if prob.is_cuda:
one_hot = one_hot.cuda()
loss = torch.masked_select(prob, one_hot)
loss = loss * reward
loss = -torch.sum(loss)
return loss
target_lstm = TargetLSTM(VOCAB_SIZE, g_emb_dim, g_hidden_dim, opt.cuda)
if opt.cuda:
target_lstm = target_lstm.cuda()
# Generate toy data using target lstm
print('Generating data ...')
generate_samples(target_lstm, BATCH_SIZE, GENERATED_NUM, POSITIVE_FILE)
# Load data from file
gen_data_iter = GenDataIter(POSITIVE_FILE, BATCH_SIZE)
original_generator = Generator(VOCAB_SIZE, g_emb_dim, g_hidden_dim, opt.cuda)
def main(generator_, discriminator, model_index):
#random.seed(SEED)
#np.random.seed(SEED)
perf_dict = {}
def main():
random.seed(SEED)
np.random.seed(SEED)
# Define Networks
generator = Generator(VOCAB_SIZE, g_emb_dim, g_hidden_dim, opt.cuda)
discriminator = Discriminator(d_num_class, VOCAB_SIZE, d_emb_dim,
d_filter_sizes, d_num_filters, d_dropout)
target_lstm = TargetLSTM(VOCAB_SIZE, g_emb_dim, g_hidden_dim,
opt.cuda) #和Generator函数只有sample方法不一样；
if opt.cuda:
generator = generator.cuda()
discriminator = discriminator.cuda()
target_lstm = target_lstm.cuda()
# Generate toy data using target lstm 也就是新建一个data，假装他就是真实数据
print('Generating data ...')
generate_samples(target_lstm, BATCH_SIZE, GENERATED_NUM,
POSITIVE_FILE) #没经过啥子训练，直接随机一个lstm产出9984*20的矩阵
# Load data from file
# 每个iter输出一个data和一个target，其中data是每个point前填个0，每个target是后面添个0
gen_data_iter = GenDataIter(POSITIVE_FILE, BATCH_SIZE)
# Pretrain Generator using MLE
gen_criterion = nn.NLLLoss(
reduction='sum'
) #You may use CrossEntropyLoss instead, if you prefer not to add an extra LogSoftmax layer .
gen_optimizer = optim.Adam(generator.parameters())
if opt.cuda:
gen_criterion = gen_criterion.cuda()
print('Pretrain with MLE ...')
for epoch in range(PRE_EPOCH_NUM): #PRE_EPOCH_NUM =120
loss = train_epoch(generator, gen_data_iter, gen_criterion,
gen_optimizer) #使得generator的参数更新，使之适应gen_data_iter
print('Epoch [%d] Model Loss: %f' % (epoch, loss)) # 9月1日
generate_samples(generator, BATCH_SIZE, GENERATED_NUM, EVAL_FILE)
eval_iter = GenDataIter(EVAL_FILE, BATCH_SIZE)
loss = eval_epoch(target_lstm, eval_iter,
gen_criterion) # generator有向target_lstm靠近吗？
print('Epoch [%d] True Loss: %f' % (epoch, loss))
# Pretrain Discriminator
dis_criterion = nn.NLLLoss(reduction='sum')
dis_optimizer = optim.Adam(discriminator.parameters())
if opt.cuda:
dis_criterion = dis_criterion.cuda()
print('Pretrain Discriminator ...')
for epoch in range(5):
generate_samples(generator, BATCH_SIZE, GENERATED_NUM, NEGATIVE_FILE)
dis_data_iter = DisDataIter(POSITIVE_FILE, NEGATIVE_FILE, BATCH_SIZE)
for _ in range(3):
loss = train_epoch(discriminator, dis_data_iter, dis_criterion,
dis_optimizer)
print('Epoch [%d], loss: %f' % (epoch, loss))
# Adversarial Training 前面的预训练的目的是什么？？？？？？
rollout = Rollout(generator, 0.8)
print('#####################################################')
print('Start Adeversatial Training...\n')
gen_gan_loss = GANLoss()
gen_gan_optm = optim.Adam(generator.parameters())
if opt.cuda:
gen_gan_loss = gen_gan_loss.cuda()
gen_criterion = nn.NLLLoss(reduction='sum')
if opt.cuda:
gen_criterion = gen_criterion.cuda()
dis_criterion = nn.NLLLoss(reduction='sum')
dis_optimizer = optim.Adam(discriminator.parameters())
if opt.cuda:
dis_criterion = dis_criterion.cuda()
for total_batch in range(TOTAL_BATCH):
## Train the generator for one step
for it in range(1):
samples = generator.sample(BATCH_SIZE, g_sequence_len)
# construct the input to the genrator, add zeros before samples and delete the last column
zeros = torch.zeros((BATCH_SIZE, 1)).type(torch.LongTensor)
if samples.is_cuda:
zeros = zeros.cuda()
inputs = Variable(
torch.cat([zeros, samples.data], dim=1)[:, :-1].contiguous())
targets = Variable(samples.data).contiguous().view(
(-1, )) # 这里我不明白，为什么inputs：（batch_size,seq_len)，而targets是一个序列
# calculate the reward
rewards = rollout.get_reward(
samples, 16, discriminator) # rewards:(batch_size,seq_len)
rewards = Variable(torch.Tensor(rewards))
rewards = torch.exp(rewards).contiguous().view((-1, ))
if opt.cuda:
rewards = rewards.cuda()
prob = generator.forward(inputs)
loss = gen_gan_loss(prob, targets, rewards)
gen_gan_optm.zero_grad()
loss.backward()
gen_gan_optm.step()
if total_batch % 1 == 0 or total_batch == TOTAL_BATCH - 1: #
generate_samples(generator, BATCH_SIZE, GENERATED_NUM, EVAL_FILE)
eval_iter = GenDataIter(EVAL_FILE, BATCH_SIZE)
loss = eval_epoch(target_lstm, eval_iter, gen_criterion)
print('Batch [%d] True Loss: %f' % (total_batch, loss))
rollout.update_params() #这里理解了的话，基本没问题了
for _ in range(4):
generate_samples(generator, BATCH_SIZE, GENERATED_NUM,
NEGATIVE_FILE)
dis_data_iter = DisDataIter(POSITIVE_FILE, NEGATIVE_FILE,
BATCH_SIZE)
for _ in range(2):
loss = train_epoch(discriminator, dis_data_iter, dis_criterion,
dis_optimizer)
def main():
random.seed(SEED)
np.random.seed(SEED)
# Build up dataset
s_train, s_test = load_from_big_file('../data/train_data_obama.txt')
# idx_to_word: List of id to word
# word_to_idx: Dictionary mapping word to id
idx_to_word, word_to_idx = fetch_vocab(s_train, s_train, s_test)
# TODO: 1. Prepare data for attention model
# input_seq, target_seq = prepare_data(DATA_GERMAN, DATA_ENGLISH, word_to_idx)
global VOCAB_SIZE
VOCAB_SIZE = len(idx_to_word)
print('VOCAB SIZE:' , VOCAB_SIZE)
# Define Networks
generator = Generator(VOCAB_SIZE, g_emb_dim, g_hidden_dim, g_sequence_len, BATCH_SIZE, opt.cuda)
discriminator = Discriminator(d_num_class, VOCAB_SIZE, d_emb_dim, d_filter_sizes, d_num_filters, d_dropout)
target_lstm = TargetLSTM(VOCAB_SIZE, g_emb_dim, g_hidden_dim, opt.cuda)
if opt.cuda:
generator = generator.cuda()
discriminator = discriminator.cuda()
target_lstm = target_lstm.cuda()
# Generate toy data using target lstm
print('Generating data ...')
generate_samples(target_lstm, BATCH_SIZE, GENERATED_NUM, POSITIVE_FILE, idx_to_word)
# Load data from file
gen_data_iter = GenDataIter(POSITIVE_FILE, BATCH_SIZE)
# Pretrain Generator using MLE
gen_criterion = nn.NLLLoss(size_average=False)
gen_optimizer = optim.Adam(generator.parameters())
if opt.cuda:
gen_criterion = gen_criterion.cuda()
print('Pretrain with MLE ...')
for epoch in range(PRE_EPOCH_NUM):
loss = train_epoch(generator, gen_data_iter, gen_criterion, gen_optimizer)
print('Epoch [%d] Model Loss: %f'% (epoch, loss))
# TODO: 2. Flags to ensure dimension of model input is handled
generate_samples(generator, BATCH_SIZE, GENERATED_NUM, EVAL_FILE)
print('Gen sampled')
eval_iter = GenDataIter(EVAL_FILE, BATCH_SIZE)
print('Iterator Done')
loss = eval_epoch(target_lstm, eval_iter, gen_criterion)
print('Epoch [%d] True Loss: %f' % (epoch, loss))
# Pretrain Discriminator
dis_criterion = nn.NLLLoss(size_average=False)
dis_optimizer = optim.Adam(discriminator.parameters())
if opt.cuda:
dis_criterion = dis_criterion.cuda()
print('Pretrain Dsicriminator ...')
for epoch in range(3):
generate_samples(generator, BATCH_SIZE, GENERATED_NUM, NEGATIVE_FILE)
dis_data_iter = DisDataIter(POSITIVE_FILE, NEGATIVE_FILE, BATCH_SIZE)
for _ in range(3):
loss = train_epoch(discriminator, dis_data_iter, dis_criterion, dis_optimizer)
print('Epoch [%d], loss: %f' % (epoch, loss))
# Adversarial Training
rollout = Rollout(generator, 0.8)
print('#####################################################')
print('Start Adeversatial Training...\n')
gen_gan_loss = GANLoss()
gen_gan_optm = optim.Adam(generator.parameters())
if opt.cuda:
gen_gan_loss = gen_gan_loss.cuda()
gen_criterion = nn.NLLLoss(size_average=False)
if opt.cuda:
gen_criterion = gen_criterion.cuda()
dis_criterion = nn.NLLLoss(size_average=False)
dis_optimizer = optim.Adam(discriminator.parameters())
if opt.cuda:
dis_criterion = dis_criterion.cuda()
for total_batch in range(TOTAL_BATCH):
## Train the generator for one step
for it in range(1):
samples = generator.sample(BATCH_SIZE, g_sequence_len)
# construct the input to the genrator, add zeros before samples and delete the last column
zeros = torch.zeros((BATCH_SIZE, 1)).type(torch.LongTensor)
if samples.is_cuda:
zeros = zeros.cuda()
inputs = Variable(torch.cat([zeros, samples.data], dim = 1)[:, :-1].contiguous())
targets = Variable(samples.data).contiguous().view((-1,))
# calculate the reward
rewards = rollout.get_reward(samples, 16, discriminator)
rewards = Variable(torch.Tensor(rewards))
if opt.cuda:
rewards = torch.exp(rewards.cuda()).contiguous().view((-1,))
prob = generator.forward(inputs)
prob = torch.reshape(prob, (prob.shape[0] * prob.shape[1], -1)) #prob.view(-1, g_emb_dim)
loss = gen_gan_loss(prob, targets, rewards)
gen_gan_optm.zero_grad()
loss.backward()
gen_gan_optm.step()
print('Batch [%d] True Loss: %f' % (total_batch, loss))
if total_batch % 10 == 0 or total_batch == TOTAL_BATCH - 1:
generate_samples(generator, BATCH_SIZE, GENERATED_NUM, EVAL_FILE)
eval_iter = GenDataIter(EVAL_FILE, BATCH_SIZE)
loss = eval_epoch(target_lstm, eval_iter, gen_criterion)
predictions = torch.max(prob, dim=1)[1]
predictions = predictions.view(BATCH_SIZE, -1)
# print('PRED SHAPE:' , predictions.shape)
for each_sen in list(predictions):
print('Sample Output:', generate_sentence_from_id(idx_to_word, each_sen, DEBUG_FILE))
sys.stdout.flush()
torch.save(generator.state_dict(), './experiment_3_10000/generator.model')
torch.save(discriminator.state_dict(), './experiment_3_10000/discriminator.model')
rollout.update_params()
for _ in range(4):
generate_samples(generator, BATCH_SIZE, GENERATED_NUM, NEGATIVE_FILE)
dis_data_iter = DisDataIter(POSITIVE_FILE, NEGATIVE_FILE, BATCH_SIZE)
for _ in range(2):
loss = train_epoch(discriminator, dis_data_iter, dis_criterion, dis_optimizer)