use adamax and 2 hidden layers for coupling

rtqichen · Jan 11, 2019 · 809f528 · 809f528
1 parent d612674
commit 809f528
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/lib/layers/coupling.py b/lib/layers/coupling.py
@@ -7,12 +7,16 @@
 class CouplingLayer(nn.Module):
     """Used in 2D experiments."""
 
-    def __init__(self, d, intermediate_dim=32, swap=False):
+    def __init__(self, d, intermediate_dim=64, swap=False):
         nn.Module.__init__(self)
         self.d = d - (d // 2)
         self.swap = swap
         self.net_s_t = nn.Sequential(
-            nn.Linear(self.d, intermediate_dim), nn.ReLU(inplace=True), nn.Linear(intermediate_dim, (d - self.d) * 2)
+            nn.Linear(self.d, intermediate_dim),
+            nn.ReLU(inplace=True),
+            nn.Linear(intermediate_dim, intermediate_dim),
+            nn.ReLU(inplace=True),
+            nn.Linear(intermediate_dim, (d - self.d) * 2),
         )
 
     def forward(self, x, logpx=None, reverse=False):

diff --git a/train_discrete_toy.py b/train_discrete_toy.py
@@ -43,7 +43,7 @@
 parser.add_argument('--JoffdiagFrobint', type=float, default=None, help="int_t ||df/dx - df_i/dx_i||_F")
 
 parser.add_argument('--save', type=str, default='experiments/cnf')
-parser.add_argument('--viz_freq', type=int, default=5000)
+parser.add_argument('--viz_freq', type=int, default=1000)
 parser.add_argument('--val_freq', type=int, default=1000)
 parser.add_argument('--log_freq', type=int, default=100)
 parser.add_argument('--gpu', type=int, default=0)
@@ -120,7 +120,7 @@ def compute_loss(args, model, batch_size=None):
     logger.info(model)
     logger.info("Number of trainable parameters: {}".format(count_parameters(model)))
 
-    optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
+    optimizer = optim.Adamax(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
 
     time_meter = utils.RunningAverageMeter(0.98)
     loss_meter = utils.RunningAverageMeter(0.98)