add free parameter and free pretrain

Author: fukuta0614

models/seqgan.py

@@ -19,6 +19,7 @@ def __init__(self, sequence_length, vocab_size, emb_dim, hidden_dim,
         self.hidden_dim = hidden_dim
         self.sequence_length = sequence_length
         self.start_token = start_token
+        self.x0 = None
         self.reward_gamma = reward_gamma
         self.g_params = []
         self.d_params = []
@@ -104,7 +105,11 @@ def decode_one_step(self, x, train=True):
             y = self.out(h)
             return y
         else:
-            h0 = self.embed(x)
+            if len(x.data.shape) == 2:
+                h0 = x
+            else:
+                h0 = self.embed(x)
+
             h = self.lstm1(h0)
             if hasattr(self, "lstm2"):
                 h = self.lstm2(h)
@@ -115,13 +120,18 @@ def decode_one_step(self, x, train=True):
             y = self.out(h)
             return y
 
-    def generate(self, batch_size, train=False, pool=None):
+    def generate(self, batch_size, train=False, pool=None, random_input=False):
         """
         :return: (batch_size, self.seq_length)
         """
 
         self.reset_state()
-        x = chainer.Variable(self.xp.asanyarray([self.start_token] * batch_size, 'int32'), volatile=True)
+        if random_input:
+            self.x0 = np.random.normal(scale=1, size=(batch_size, self.emb_dim))
+            x = chainer.Variable(self.xp.asanyarray(self.x0, 'float32'), volatile=True)
+        else:
+            x = chainer.Variable(self.xp.asanyarray([self.start_token] * batch_size, 'int32'), volatile=True)
+
         gen_x = np.zeros((batch_size, self.sequence_length), 'int32')
 
         for i in range(self.sequence_length):
@@ -169,7 +179,7 @@ def pretrain_step(self, x_input):
 
         return accum_loss / self.sequence_length
 
-    def reinforcement_step(self, x_input, rewards, g_steps):
+    def reinforcement_step(self, x_input, rewards, g_steps, random_input=False):
         """
         :param x_input: (batch_size, seq_length)
         :param rewards: (batch_size, seq_length)
@@ -181,7 +191,10 @@ def reinforcement_step(self, x_input, rewards, g_steps):
         accum_loss = 0
         for j in range(self.sequence_length):
             if j == 0:
-                x = chainer.Variable(self.xp.asanyarray([self.start_token] * batch_size, 'int32'))
+                if random_input:
+                    x = chainer.Variable(self.xp.asanyarray(self.x0, 'float32'))
+                else:
+                    x = chainer.Variable(self.xp.asanyarray([self.start_token] * batch_size, 'int32'))
             else:
                 x = chainer.Variable(self.xp.asanyarray(x_input[:, j - 1], 'int32'))
 

models/text_cnn.py

@@ -39,7 +39,12 @@ def __init__(self, num_classes, vocab_size,
 
     def forward(self, x_input, ratio=0.5, train=True):
 
-        batch_size, seq_length = x_input.shape
+        try:
+            batch_size, seq_length = x_input.shape
+        except:
+            batch_size = len(x_input)
+            seq_length = len(x_input[0])
+
         x = chainer.Variable(self.xp.asarray(x_input, 'int32'))
 
         # embedding

oracle_test/run_sequence_gan.py

@@ -153,6 +153,7 @@ def significance_test(session, target_lstm, data_loader, output_file):
 # generator
 generator = SeqGAN(seq_length, vocab_size, gen_emb_dim, gen_hidden_dim, start_token, oracle=True).to_gpu()
 if args.gen:
+    print(args.gen)
     serializers.load_hdf5(args.gen, generator)
 
 # discriminator
@@ -212,10 +213,9 @@ def significance_test(session, target_lstm, data_loader, output_file):
         summary = sess.run(target_loss_summary, feed_dict={loss_: test_loss})
         summary_writer.add_summary(summary, test_count)
 
-    with open(os.path.join(out_dir, "models", "gen_pretrain.model"), "wb") as f:
-        pickle.dump(generator, f)
-    with open(os.path.join(out_dir, "models", "gen_pretrain.opt"), 'wb') as f:
-        pickle.dump(gen_optimizer, f)
+    serializers.save_hdf5(os.path.join(out_dir, "models", "gen_pretrain.model"), generator)
+    serializers.save_hdf5(os.path.join(out_dir, "models", "gen_pretrain.opt"), gen_optimizer)
+
 else:
     test_count = gen_pretrain_epoch
     test_loss = generator.target_loss(target_lstm, 1000, gen_batch_size, sess)
@@ -258,8 +258,8 @@ def significance_test(session, target_lstm, data_loader, output_file):
         summary_writer.add_summary(summary, dis_train_count)
         summary = sess.run(dis_acc_summary, feed_dict={loss_: np.mean(sum_train_accuracy)})
         summary_writer.add_summary(summary, dis_train_count)
-    with open(os.path.join(out_dir, "models", "dis_pretrain.model"), "wb") as f:
-        pickle.dump(discriminator, f)
+    serializers.save_hdf5(os.path.join(out_dir, "models", "dis_pretrain.model"), discriminator)
+    serializers.save_hdf5(os.path.join(out_dir, "models", "dis_pretrain.opt"), dis_optimizer)
 
 # roll out generator
 rollout_generator = copy.deepcopy(generator)
@@ -275,10 +275,10 @@ def significance_test(session, target_lstm, data_loader, output_file):
     print('total batch: ', epoch)
 
     for step in range(g_steps):
-        samples = generator.generate(gen_batch_size, train=True)
+        samples = generator.generate(gen_batch_size, train=True, random_input=True)
         rewards = rollout_generator.get_rewards(samples, discriminator, rollout_num=16, pool=pool, gpu=args.gpu)
         print(rewards[:30])
-        loss = generator.reinforcement_step(samples, rewards, g_steps=g_steps)
+        loss = generator.reinforcement_step(samples, rewards, g_steps=g_steps, random_input=True)
         gen_optimizer.zero_grads()
         loss.backward()
         gen_optimizer.update()