def LSTM(): #Step 1: Read data from files and put them into list test_sentences = list() with open(args.test_data, 'r') as f: for line in f: test_sentences.append(line.replace("<eos>", "")) temp = list() for i in range(0, len(test_sentences)): if test_sentences[i] != '\n': temp.append(test_sentences[i]) test_sentences = temp print(test_sentences) real_sentences = list() with codecs.open(args.real_data,'r',encoding='utf8',errors='ignore') as f: for line in f: real_sentences.append(line) #Step 2: BLEU Score print("LSTM - LSTM double layer encoding") for i in range(1, args.gram + 1): bleu = BLEU(test_sentences, real_sentences, i) bleu_score = bleu.get_score(ignore=False) print("BLEU{} score:{}".format(i,bleu_score))
def leakGAN(): test_sentences = list() with codecs.open(args.test_data, 'r',encoding='utf8',errors='ignore') as f: for line in f: line = line.split(' ', 1)[1] #remove initial EOS test_sentences.append(line) real_sentences = list() with codecs.open(args.real_data,'r',encoding='utf8',errors='ignore') as f: for line in f: real_sentences.append(line) #Step 2: BLEU Score print("LSTM - LeakGAN double layer encoding") for i in range(1, args.gram + 1): bleu = BLEU(test_sentences, real_sentences, i) bleu_score = bleu.get_score(ignore=False) print("BLEU{} score:{}".format(i,bleu_score))
class LeakGANInstructor(BasicInstructor): 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)) 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 _run(self): for inter_num in range(cfg.inter_epoch): self.log.info('>>> Interleaved Round %d...' % inter_num) self.sig.update() # update signal if self.sig.pre_sig: # =====DISCRIMINATOR PRE-TRAINING===== if not cfg.dis_pretrain: self.log.info('Starting Discriminator Training...') self.train_discriminator(cfg.d_step, cfg.d_epoch) if cfg.if_save and not cfg.if_test: torch.save(self.dis.state_dict(), cfg.pretrained_dis_path) print('Save pre-trained discriminator: {}'.format( cfg.pretrained_dis_path)) # =====GENERATOR MLE TRAINING===== if not cfg.gen_pretrain: self.log.info('Starting Generator MLE Training...') self.pretrain_generator(cfg.MLE_train_epoch) if cfg.if_save and not cfg.if_test: torch.save(self.gen.state_dict(), cfg.pretrained_gen_path) print('Save pre-trained generator: {}'.format( cfg.pretrained_gen_path)) else: self.log.info( '>>> Stop by pre_signal! Skip to adversarial training...') break # =====ADVERSARIAL TRAINING===== self.log.info('Starting Adversarial Training...') self.log.info('Initial generator: %s' % (str(self.cal_metrics(fmt_str=True)))) for adv_epoch in range(cfg.ADV_train_epoch): self.log.info('-----\nADV EPOCH %d\n-----' % adv_epoch) self.sig.update() if self.sig.adv_sig: self.adv_train_generator(cfg.ADV_g_step) # Generator self.train_discriminator(cfg.ADV_d_step, cfg.ADV_d_epoch, 'ADV') # Discriminator if adv_epoch % cfg.adv_log_step == 0: if cfg.if_save and not cfg.if_test: self._save('ADV', adv_epoch) else: self.log.info( '>>> Stop by adv_signal! Finishing adversarial training...' ) break def _test(self): print('>>> Begin test...') self._run() pass def pretrain_generator(self, epochs): """ Max Likelihood Pretraining for the gen - gen_opt: [mana_opt, work_opt] """ for epoch in range(epochs): self.sig.update() if self.sig.pre_sig: pre_mana_loss = 0 pre_work_loss = 0 # =====Train===== for i, data in enumerate(self.oracle_data.loader): inp, target = data['input'], data['target'] if cfg.CUDA: inp, target = inp.cuda(), target.cuda() mana_loss, work_loss = self.gen.pretrain_loss( target, self.dis) self.optimize_multi(self.gen_opt, [mana_loss, work_loss]) pre_mana_loss += mana_loss.data.item() pre_work_loss += work_loss.data.item() pre_mana_loss = pre_mana_loss / len(self.oracle_data.loader) pre_work_loss = pre_work_loss / len(self.oracle_data.loader) # =====Test===== if epoch % cfg.pre_log_step == 0: self.log.info( '[MLE-GEN] epoch %d : pre_mana_loss = %.4f, pre_work_loss = %.4f, %s' % (epoch, pre_mana_loss, pre_work_loss, self.cal_metrics(fmt_str=True))) if cfg.if_save and not cfg.if_test: self._save('MLE', epoch) else: self.log.info( '>>> Stop by pre signal, skip to adversarial training...') break 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 = self.gen_data.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 train_discriminator(self, d_step, d_epoch, phrase='MLE'): """ Training the discriminator on real_data_samples (positive) and generated samples from gen (negative). Samples are drawn d_step times, and the discriminator is trained for d_epoch d_epoch. """ for step in range(d_step): # prepare loader for training pos_samples = self.oracle_data.target neg_samples = self.gen.sample(cfg.samples_num, cfg.batch_size, self.dis) self.dis_data.reset(pos_samples, neg_samples) for epoch in range(d_epoch): # =====Train===== d_loss, train_acc = self.train_dis_epoch( self.dis, self.dis_data.loader, self.dis_criterion, self.dis_opt) # =====Test===== self.log.info( '[%s-DIS] d_step %d: d_loss = %.4f, train_acc = %.4f,' % (phrase, step, d_loss, train_acc)) def cal_metrics(self, fmt_str=False): self.gen_data.reset( self.gen.sample(cfg.samples_num, cfg.batch_size, self.dis)) self.bleu3.test_text = tensor_to_tokens(self.gen_data.target, self.index_word_dict) bleu3_score = self.bleu3.get_score(ignore=False) with torch.no_grad(): gen_nll = 0 for data in self.oracle_data.loader: inp, target = data['input'], data['target'] if cfg.CUDA: inp, target = inp.cuda(), target.cuda() loss = self.gen.batchNLLLoss(target, self.dis) gen_nll += loss.item() gen_nll /= len(self.oracle_data.loader) if fmt_str: return 'BLEU-3 = %.4f, gen_NLL = %.4f,' % (bleu3_score, gen_nll) return bleu3_score, gen_nll def _save(self, phrase, epoch): torch.save( self.gen.state_dict(), cfg.save_model_root + 'gen_{}_{:05d}.pt'.format(phrase, epoch)) save_sample_path = cfg.save_samples_root + 'samples_{}_{:05d}.txt'.format( phrase, epoch) samples = self.gen.sample(cfg.batch_size, cfg.batch_size, self.dis) write_tokens(save_sample_path, tensor_to_tokens(samples, self.index_word_dict))
class Model(): def __init__(self, encoder_layer_num, decoder_layer_num, hidden_dim, batch_size, learning_rate, dropout, init_train = True): self.encoder_layer_num = encoder_layer_num self.decoder_layer_num = decoder_layer_num self.hidden_dim = hidden_dim self.batch_size = batch_size self.learning_rate = learning_rate self.dropout = dropout self.init_train = init_train #---------fix---------- self.vocab_size = cfg.vocab_size self.max_length = cfg.max_length self.embedding_matrix = make_embedding_matrix(cfg.all_captions) self.SOS_token = cfg.SOS_token self.EOS_token = cfg.EOS_token self.idx2word_dict = load_dict() #---------------------- self.bleu = BLEU('BLEU', gram=[2,3,4,5]) #self.bleu.reset(test_text = gen_tokens, real_text = self.test_data.tokens) if init_train: self._init_train() train_week_stock, train_month_stock, t_month_stock,train_input_cap_vector, train_output_cap_vector = load_training_data() self.train_data = batch_generator(train_week_stock, train_month_stock, t_month_stock,train_input_cap_vector, train_output_cap_vector, self.batch_size) self.total_iter = len(train_input_cap_vector) self._init_eval() val_week_stock, val_month_stock, val_t_month_stock,val_input_cap_vector, val_output_cap_vector = load_val_data() self.val_data = batch_generator(val_week_stock, val_month_stock, val_t_month_stock,val_input_cap_vector, val_output_cap_vector, self.batch_size) self.val_total_iter = len(val_input_cap_vector) # gpu 탄력적으로 사용. def gpu_session_config(self): config = tf.ConfigProto() config.gpu_options.allow_growth = True return config def _init_train(self): self.train_graph = tf.Graph() with self.train_graph.as_default(): with tf.variable_scope('encoder_input'): self.week_input = tf.placeholder(tf.float64, shape= [None, 7], name='week_input') self.month_input = tf.placeholder(tf.float64, shape=[None, 28], name='month_input') self.t_month_input = tf.placeholder(tf.float64, shape=[None, 84], name='t_month_input') with tf.variable_scope("decoder_input"): self.decoder_input = tf.placeholder(tf.int32, [None, self.max_length], name='input') self.decoder_target = tf.placeholder(tf.int32, [None, self.max_length], name='target') self.decoder_targets_length = tf.placeholder(tf.int32, shape = [self.batch_size, ], name = 'targets_length') encoded_output, encoded_state = encoder_module(self.week_input, self.month_input, self.t_month_input, self.encoder_layer_num, self.decoder_layer_num, self.hidden_dim) decoder_output, decoder_state = decoder_module(encoded_state, encoded_output, self.decoder_input, self.decoder_targets_length, self.embedding_matrix, self.decoder_layer_num, self.hidden_dim, self.max_length, self.vocab_size, self.batch_size, self.dropout, self.SOS_token, self.EOS_token, train = True) self.logits = decoder_output.rnn_output # traning output self.sample_id = decoder_output.sample_id self._init_optimizer() self.train_init = tf.global_variables_initializer() self.train_saver = tf.train.Saver() self.train_session = tf.Session(graph=self.train_graph, config = self.gpu_session_config()) def _init_optimizer(self): #loss mask mask = tf.cast(tf.sequence_mask(self.decoder_targets_length, self.max_length),tf.float64) self.loss = tf.contrib.seq2seq.sequence_loss(logits= self.logits, targets = self.decoder_target, weights = mask, average_across_timesteps = True, average_across_batch = True) #tf.summary.scalar('loss', self.loss) #self.summary_op = tf.summary.merge_all() params = tf.trainable_variables() gradients = tf.gradients(self.loss, params) clipped_gradients, _ = tf.clip_by_global_norm(gradients,5.0) self.optimizer = tf.train.AdamOptimizer(self.learning_rate).apply_gradients(zip(clipped_gradients, params)) # batch 단위로 계산 def cal_metrics(self, infer_text, real_text): self.bleu.reset(infer_text = infer_text, real_text = real_text) return self.bleu.get_score() # bleu, greedy/beam search init def _init_eval(self): self.eval_graph = tf.Graph() with self.eval_graph.as_default(): self.eval_week_input = tf.placeholder(tf.float64, shape= [None, 7]) self.eval_month_input = tf.placeholder(tf.float64, shape=[None, 28]) self.eval_t_month_input = tf.placeholder(tf.float64, shape=[None, 84]) self.eval_decoder_targets_length = tf.placeholder(tf.int32, shape = [self.batch_size, ]) eval_encoded_output, eval_encoded_state = encoder_module(self.eval_week_input, self.eval_month_input, self.eval_t_month_input, self.encoder_layer_num, self.decoder_layer_num, self.hidden_dim) self.eval_decoder_output, eval_decoder_state = decoder_module(eval_encoded_state, eval_encoded_output, None, self.eval_decoder_targets_length, self.embedding_matrix, self.decoder_layer_num, self.hidden_dim, self.max_length, self.vocab_size, self.batch_size, self.dropout, self.SOS_token, self.EOS_token, train = False) self.predicted_ids = tf.identity(self.eval_decoder_output.predicted_ids) self.eval_saver = tf.train.Saver() self.eval_session = tf.Session(graph=self.eval_graph,config=self.gpu_session_config()) def train_epoch(self, epochs): if not self.init_train: raise Exception('Train graph is not inited') with self.train_graph.as_default(): if os.path.isfile(cfg.save_path + '.meta'): print("##########################") print('# Model restore.. #') print("##########################") self.train_saver.restore(self.train_session, cfg.save_path) else: self.train_session.run(self.train_init) total_loss = 0 total_step = 0 start_time =time.time() for e in range(epochs): for step in range(self.total_iter// self.batch_size): data = next(self.train_data) week_stock = data['week_stock'] month_stock = data['month_stock'] t_month_stock = data['t_month_stock'] decoder_input = data['decoder_input'] decoder_target = data['decoder_target'] batch_seq = batch_seq_len(data['decoder_target']) _, loss, sample_id = self.train_session.run([self.optimizer, self.loss, self.sample_id], feed_dict = {self.week_input : week_stock, self.month_input : month_stock, self.t_month_input : t_month_stock, self.decoder_input : decoder_input, self.decoder_target : decoder_target, self.decoder_targets_length : batch_seq}) # total_loss += loss # total_step += self.total_iter # loss = total_loss/total_step end = time.time() print('epoch: {}|{} minibatch loss: {:.6f} Time: {:.1f} min'.format(e+1, epochs, loss, (end-start_time)/60 )) if e % 50 ==0: self.train_saver.save(self.train_session, cfg.save_path) #랜덤 sid 선택, training output_text sid = random.randint(0, self.batch_size-1) target_text = decode_text(decoder_target[sid],self.idx2word_dict) output_text = decode_text(sample_id[sid],self.idx2word_dict) print('============ training sample text =============') print('training_target :' + target_text) print('training_output :' + output_text) print('===============================================') self.eval() def eval(self): with self.eval_graph.as_default(): self.eval_saver.restore(self.eval_session, cfg.save_path) all_bleu = [0] * 4 eval_mask_weights = tf.ones(shape=[self.batch_size, self.max_length],dtype=tf.float64) for step in range(self.val_total_iter//self.batch_size): data = next(self.val_data) week_stock = data['week_stock'] month_stock = data['month_stock'] t_month_stock = data['t_month_stock'] batch_seq = batch_seq_len(data['decoder_target']) #beam search_output beam_output = self.eval_session.run([self.predicted_ids], feed_dict = {self.eval_week_input : week_stock, self.eval_month_input : month_stock, self.eval_t_month_input : t_month_stock, self.eval_decoder_targets_length : batch_seq }) target_text = idx_to_text(data['decoder_input'][:,1:],self.idx2word_dict) target_text = remove_sent_pad(target_text) beam_output = np.squeeze(np.array(beam_output),axis=0) output_text = idx_to_text(beam_output[:,:,0], self.idx2word_dict) bleu_score = self.cal_metrics(target_text, output_text) for idx,score in enumerate(bleu_score): all_bleu[idx] += score print('================ BLEU score ================') for idx, bleu in enumerate(bleu_score):#2,3,4,5 print('BLEU-{} : {}'.format(idx+2, bleu)) sid = random.randint(0, self.batch_size-1) target_text = decode_text(data['decoder_target'][sid],self.idx2word_dict) output_text = decode_text(beam_output[sid,:,0],self.idx2word_dict) print('============= Beam search text =============') print('infer_target : ' + target_text) print('beam_search : ' + output_text) print('============================================')
class LeakGANInstructor(BasicInstructor): 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 _run(self): for inter_num in range(cfg.inter_epoch): self.log.info('>>> Interleaved Round %d...' % inter_num) self.sig.update() # update signal if self.sig.pre_sig: # =====DISCRIMINATOR PRE-TRAINING===== if not cfg.dis_pretrain: self.log.info('Starting Discriminator Training...') self.train_discriminator(cfg.d_step, cfg.d_epoch) if cfg.if_save and not cfg.if_test: torch.save(self.dis.state_dict(), cfg.pretrained_dis_path) print('Save pre-trained discriminator: {}'.format( cfg.pretrained_dis_path)) # =====GENERATOR MLE TRAINING===== if not cfg.gen_pretrain: self.log.info('Starting Generator MLE Training...') self.pretrain_generator(cfg.MLE_train_epoch) if cfg.if_save and not cfg.if_test: torch.save(self.gen.state_dict(), cfg.pretrained_gen_path) print('Save pre-trained generator: {}'.format( cfg.pretrained_gen_path)) else: self.log.info( '>>> Stop by pre_signal! Skip to adversarial training...') break # =====ADVERSARIAL TRAINING===== self.log.info('Starting Adversarial Training...') self.log.info('Initial generator: %s' % (str(self.cal_metrics(fmt_str=True)))) for adv_epoch in range(cfg.ADV_train_epoch): self.log.info('-----\nADV EPOCH %d\n-----' % adv_epoch) self.sig.update() if self.sig.adv_sig: self.adv_train_generator(cfg.ADV_g_step) # Generator self.train_discriminator(cfg.ADV_d_step, cfg.ADV_d_epoch, 'ADV') # Discriminator if adv_epoch % cfg.adv_log_step == 0: if cfg.if_save and not cfg.if_test: self._save('ADV', adv_epoch) else: self.log.info( '>>> Stop by adv_signal! Finishing adversarial training...' ) break def string2bins(self, bit_string, n_bins): n_bits = int(math.log(n_bins, 2)) return [ bit_string[i:i + n_bits] for i in range(0, len(bit_string), n_bits) ] def LSTM_layer_1(self, intermediate_file, bins_num): print('>>> Begin test...') print('Begin with LSTM Layer') #First layer- LSTM layer epoch_start_time = time.time() seed = 1111 data_root = './decode/' #Reproducibility torch.manual_seed(seed) if cfg.CUDA: torch.cuda.manual_seed(seed) with open("leakGAN_instructor/real_data/emnlp_news.pt", 'rb') as f: self.lstm = torch.load(f) if cfg.CUDA: self.lstm.cuda() emnlp_data = 'dataset/emnlp_news/' corpus = dataa.Corpus(emnlp_data) ntokens = len(corpus.dictionary) idx2word_file = data_root + "idx2word_1.txt" word2idx_file = data_root + "word2idx_1.txt" with open(idx2word_file, "wb") as fp: #Pickling pickle.dump(corpus.dictionary.idx2word, fp) with open(word2idx_file, "wb") as fp: #Pickling pickle.dump(corpus.dictionary.word2idx, fp) hidden = self.lstm.init_hidden(1) input = torch.randint(ntokens, (1, 1), dtype=torch.long) if cfg.CUDA: input.data = input.data.cuda() print("Finished Initializing LSTM Model") #Step 1: Get secret data secret_file = open("leakGAN_instructor/real_data/secret_file.txt", 'r') secret_data = secret_file.read() #Step 2: Compress string into binary string bit_string = ''.join( bin(ord(letter))[2:].zfill(8) for letter in secret_data) #print(bit_string) bit_string = '111011100101000111000011110111101111110111000110011010110110' #In the first step we will use 256 bins (8 bit representation each) to convert so that we can convert 64 bits into 8 word #Step 3: Divide into bins secret_text = [ int(i, 2) for i in self.string2bins(bit_string, bins_num) ] #convert to bins #Step 4: Divide vocabulary into bins, zero words not in the bin if bins_num >= 2: tokens = list(range(ntokens)) #indecies of words random.shuffle(tokens) #randomize #Words in each bin words_in_bin = int(ntokens / bins_num) #leftovers should be also included in the leftover = int(ntokens % bins_num) bins = [ tokens[i:i + words_in_bin] for i in range(0, ntokens - leftover, words_in_bin) ] # words to keep in each bin for i in range(len(bins)): if (i == leftover): break bins[i].append(tokens[i + words_in_bin * bins_num]) print("Len of bins in 1st layer: {}".format(len(bins))) #save bins into key 1 key1 = data_root + "lstm_key1.txt" with open(key1, "wb") as fp: #Pickling pickle.dump(bins, fp) zero = [list(set(tokens) - set(bin_)) for bin_ in bins] print('Finished Initializing First LSTM Layer') print('time: {:5.2f}s'.format(time.time() - epoch_start_time)) print('-' * 89) intermediate_file = data_root + intermediate_file with open(intermediate_file, 'w') as outf: w = 0 i = 1 temperature = 1.5 bin_sequence_length = len(secret_text[:]) # 85 print("bin sequence length", bin_sequence_length) #32 while i <= bin_sequence_length: epoch_start_time = time.time() output, hidden = self.lstm(input, hidden) zero_index = zero[secret_text[:][i - 1]] zero_index = torch.LongTensor(zero_index) word_weights = output.squeeze().data.div( temperature).exp().cpu() word_weights.index_fill_(0, zero_index, 0) word_idx = torch.multinomial(word_weights, 1)[0] input.data.fill_(word_idx) word = corpus.dictionary.idx2word[word_idx] i += 1 w += 1 word = word.encode('ascii', 'ignore').decode('ascii') outf.write(word + ' ') print("Generated intermediate short steganographic text") print("Intermediate text saved in following file: {}".format( intermediate_file)) def LSTM_layer_2(self, secret_file, final_file, bins_num): print('Final LSTM Layer') #First layer- LSTM layer data_root = './decode/' epoch_start_time = time.time() seed = 1111 #Reproducibility torch.manual_seed(seed) if cfg.CUDA: torch.cuda.manual_seed(seed) with open("leakGAN_instructor/real_data/emnlp_news.pt", 'rb') as f: self.lstm = torch.load(f) if cfg.CUDA: self.lstm.cuda() emnlp_data = 'dataset/emnlp_news/' corpus = dataa.Corpus(emnlp_data) #save dictionary idx2word_file = data_root + "idx2word_2.txt" word2idx_file = data_root + "word2idx_2.txt" with open(idx2word_file, "wb") as fp: #Pickling pickle.dump(corpus.dictionary.idx2word, fp) with open(word2idx_file, "wb") as fp: #Pickling pickle.dump(corpus.dictionary.word2idx, fp) ntokens = len(corpus.dictionary) hidden = self.lstm.init_hidden(1) input = torch.randint(ntokens, (1, 1), dtype=torch.long) if cfg.CUDA: input.data = input.data.cuda() print("Finished Initializing LSTM Model") #Step 1: Get secret data secret_file = open(data_root + secret_file, 'r') secret_data = secret_file.read().split() #Step 2: Compress string into binary string bit_string = '' for data in secret_data: print("Data: {}".format(data)) idWord = corpus.dictionary.word2idx[data] bit_string += '{0:{fill}13b}'.format(int(idWord), fill='0') #print(ntokens) print("Bit String: {}".format(bit_string)) print("Length of Bit String: {}".format(len(bit_string))) #print(bit_string) #bit_string = '111011100101000111000011110111101111110111000110011010110110' #In the first step we will use 256 bins (8 bit representation each) to convert so that we can convert 64 bits into 8 word #Step 3: Divide into bins secret_text = [ int(i, 2) for i in self.string2bins(bit_string, bins_num) ] #convert to bins #Step 4: Divide vocabulary into bins, zero words not in the bin if bins_num >= 2: tokens = list(range(ntokens)) #indecies of words random.shuffle(tokens) #randomize #Words in each bin words_in_bin = int(ntokens / bins_num) #leftovers should be also included in the leftover = int(ntokens % bins_num) bins = [ tokens[i:i + words_in_bin] for i in range(0, ntokens - leftover, words_in_bin) ] # words to keep in each bin for i in range(0, leftover): bins[i].append(tokens[i + words_in_bin * bins_num]) #save bins into key 1 key1 = data_root + "lstm_key2.txt" with open(key1, "wb") as fp: #Pickling pickle.dump(bins, fp) zero = [list(set(tokens) - set(bin_)) for bin_ in bins] print('Finished Initializing Second LSTM Layer') print('time: {:5.2f}s'.format(time.time() - epoch_start_time)) print('-' * 89) final_file = data_root + final_file with open(final_file, 'w') as outf: w = 0 i = 1 temperature = 1.5 bin_sequence_length = len(secret_text[:]) # 85 print("bin sequence length", bin_sequence_length) #32 while i <= bin_sequence_length: epoch_start_time = time.time() output, hidden = self.lstm(input, hidden) zero_index = zero[secret_text[:][i - 1]] zero_index = torch.LongTensor(zero_index) word_weights = output.squeeze().data.div( temperature).exp().cpu() word_weights.index_fill_(0, zero_index, 0) word_idx = torch.multinomial(word_weights, 1)[0] input.data.fill_(word_idx) word = corpus.dictionary.idx2word[word_idx] i += 1 w += 1 word = word.encode('ascii', 'ignore').decode('ascii') outf.write(word + ' ') print("Generated final steganographic text") print("Final text saved in following file: {}".format( str(data_root + final_file))) def leakGAN_layer(self, secret_file, final_file, bins_num): #Second Layer = LeakGAN layer print('>>> Begin Second Layer...') data_root = './decode/' torch.nn.Module.dump_patches = True epoch_start_time = time.time() # Set the random seed manually for reproducibility. seed = 1111 #Step 1: load the most accurate model with open("leakGAN_instructor/real_data/gen_ADV_00028.pt", 'rb') as f: self.gen.load_state_dict(torch.load(f)) print("Finish Loading") self.gen.eval() #Step 1: Get Intermediate text secret_file = data_root + secret_file secret_file = open(secret_file, 'r') secret_data = secret_file.read().split() #Step 2: Compress string into binary string bit_string = '' #You need LSTM Corpus for that emnlp_data = 'dataset/emnlp_news/' corpus = dataa.Corpus(emnlp_data) for data in secret_data: print("Data: {}".format(data)) idWord = corpus.dictionary.word2idx[data] bit_string += '{0:{fill}13b}'.format(int(idWord), fill='0') secret_text = [ int(i, 2) for i in self.string2bins(bit_string, bins_num) ] #convert to bins corpus_leak = self.index_word_dict if bins_num >= 2: ntokens = len(corpus_leak) tokens = list(range(ntokens)) # * args.replication_factor #print(ntokens) random.shuffle(tokens) #Words in each bin words_in_bin = int(ntokens / bins_num) #leftovers should be also included in the leftover = int(ntokens % bins_num) bins = [ tokens[i:i + words_in_bin] for i in range(0, ntokens - leftover, words_in_bin) ] # words to keep in each bin for i in range(0, leftover): bins[i].append(tokens[i + words_in_bin * bins_num]) #save bins into leakGAN key key2 = data_root + 'leakGAN_key.txt' with open(key2, "wb") as fp: #Pickling pickle.dump(bins, fp) zero = [list(set(tokens) - set(bin_)) for bin_ in bins] print('Finished Initializing Second LeakGAN Layer') print('time: {:5.2f}s'.format(time.time() - epoch_start_time)) print('-' * 89) out_file = data_root + final_file w = 0 i = 1 bin_sequence_length = len(secret_text[:]) print("bin sequence length", bin_sequence_length) batch_size = cfg.batch_size seq_len = cfg.max_seq_len feature_array = torch.zeros( (batch_size, seq_len + 1, self.gen.goal_out_size)) goal_array = torch.zeros( (batch_size, seq_len + 1, self.gen.goal_out_size)) leak_out_array = torch.zeros((batch_size, seq_len + 1, cfg.vocab_size)) samples = torch.zeros(batch_size, seq_len + 1).long() work_hidden = self.gen.init_hidden(batch_size) mana_hidden = self.gen.init_hidden(batch_size) leak_inp = torch.LongTensor([cfg.start_letter] * batch_size) real_goal = self.gen.goal_init[:batch_size, :] if cfg.CUDA: feature_array = feature_array.cuda() goal_array = goal_array.cuda() leak_out_array = leak_out_array.cuda() goal_array[:, 0, :] = real_goal # g0 = goal_init if_sample = True no_log = False index = cfg.start_letter while i <= seq_len: dis_inp = torch.zeros(batch_size, bin_sequence_length).long() if i > 1: dis_inp[:, :i - 1] = samples[:, :i - 1] # cut sentences leak_inp = samples[:, i - 2] if torch.cuda.is_available(): dis_inp = dis_inp.cuda() leak_inp = leak_inp.cuda() feature = self.dis.get_feature(dis_inp).unsqueeze(0) #print(feature) feature_array[:, i - 1, :] = feature.squeeze(0) out, cur_goal, work_hidden, mana_hidden = self.gen(index, leak_inp, work_hidden, mana_hidden, feature, real_goal, no_log=no_log, train=False) leak_out_array[:, i - 1, :] = out goal_array[:, i, :] = cur_goal.squeeze(1) if i > 0 and i % self.gen.step_size == 0: real_goal = torch.sum(goal_array[:, i - 3:i + 1, :], dim=1) if i / self.gen.step_size == 1: real_goal += self.gen.goal_init[:batch_size, :] # Sample one token if not no_log: out = torch.exp(out) zero_index = zero[secret_text[:][ i - 1]] #indecies that has to be zeroed, as they are not in the current bin #zero_index.append(0) zero_index = torch.LongTensor(zero_index) if cfg.CUDA: zero_index = zero_index.cuda() temperature = 1.5 word_weights = out word_weights = word_weights.index_fill_( 1, zero_index, 0) #make all the indecies zero if they are not in the bin word_weights = torch.multinomial(word_weights, 1).view( -1) #choose one word with highest probability for each sample #print("Out after: {}".format(word_weights)) samples[:, i] = word_weights leak_inp = word_weights i += 1 w += 1 leak_out_array = leak_out_array[:, :seq_len, :] tokens = [] write_tokens(out_file, tensor_to_tokens(samples, self.index_word_dict)) print("Generated final steganographic text") print("Final steganographic text saved in following file: {}".format( out_file)) def _test_2_layers(self): self.LSTM_layer_1("intermediate.txt", 4096) if cfg.leakGAN: self.leakGAN_layer("intermediate.txt", "final_leakgan.txt", 4) else: self.LSTM_layer_2("intermediate.txt", "final_lstm.txt", 4) def _test(self): print('>>> Begin test...') def pretrain_generator(self, epochs): """ Max Likelihood Pretraining for the gen - gen_opt: [mana_opt, work_opt] """ for epoch in range(epochs): self.sig.update() if self.sig.pre_sig: pre_mana_loss = 0 pre_work_loss = 0 # =====Train===== for i, data in enumerate(self.oracle_data.loader): inp, target = data['input'], data['target'] if cfg.CUDA: inp, target = inp.cuda(), target.cuda() mana_loss, work_loss = self.gen.pretrain_loss( target, self.dis) self.optimize_multi(self.gen_opt, [mana_loss, work_loss]) pre_mana_loss += mana_loss.data.item() pre_work_loss += work_loss.data.item() pre_mana_loss = pre_mana_loss / len(self.oracle_data.loader) pre_work_loss = pre_work_loss / len(self.oracle_data.loader) # =====Test===== if epoch % cfg.pre_log_step == 0: self.log.info( '[MLE-GEN] epoch %d : pre_mana_loss = %.4f, pre_work_loss = %.4f, %s' % (epoch, pre_mana_loss, pre_work_loss, self.cal_metrics(fmt_str=True))) if cfg.if_save and not cfg.if_test: self._save('MLE', epoch) else: self.log.info( '>>> Stop by pre signal, skip to adversarial training...') break 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 = self.gen_data.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 train_discriminator(self, d_step, d_epoch, phrase='MLE'): """ Training the discriminator on real_data_samples (positive) and generated samples from gen (negative). Samples are drawn d_step times, and the discriminator is trained for d_epoch d_epoch. """ for step in range(d_step): # prepare loader for training pos_samples = self.oracle_data.target neg_samples = self.gen.sample(cfg.samples_num, cfg.batch_size, self.dis) self.dis_data.reset(pos_samples, neg_samples) for epoch in range(d_epoch): # =====Train===== d_loss, train_acc = self.train_dis_epoch( self.dis, self.dis_data.loader, self.dis_criterion, self.dis_opt) # =====Test===== self.log.info( '[%s-DIS] d_step %d: d_loss = %.4f, train_acc = %.4f,' % (phrase, step, d_loss, train_acc)) def cal_metrics(self, fmt_str=False): self.gen_data.reset( self.gen.sample(cfg.samples_num, cfg.batch_size, self.dis)) self.bleu3.test_text = tensor_to_tokens(self.gen_data.target, self.index_word_dict) bleu3_score = self.bleu3.get_score(ignore=False) with torch.no_grad(): gen_nll = 0 for data in self.oracle_data.loader: inp, target = data['input'], data['target'] if cfg.CUDA: inp, target = inp.cuda(), target.cuda() loss = self.gen.batchNLLLoss(target, self.dis) gen_nll += loss.item() gen_nll /= len(self.oracle_data.loader) if fmt_str: return 'BLEU-3 = %.4f, gen_NLL = %.4f,' % (bleu3_score, gen_nll) return bleu3_score, gen_nll def _save(self, phrase, epoch): torch.save( self.gen.state_dict(), cfg.save_model_root + 'gen_{}_{:05d}.pt'.format(phrase, epoch)) save_sample_path = cfg.save_samples_root + 'samples_{}_{:05d}.txt'.format( phrase, epoch) samples = self.gen.sample(cfg.batch_size, cfg.batch_size, self.dis) write_tokens(save_sample_path, tensor_to_tokens(samples, self.index_word_dict))