lr mult

JesseYang · Apr 15, 2018 · c6cab69 · c6cab69
1 parent 69f23ff
commit c6cab69
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 14 deletions.
diff --git a/cfgs/config_normal.py b/cfgs/config_normal.py
@@ -22,7 +22,8 @@
 _.skip_adj = False
 
 _.augmentation = True
-_.backbone_grad_scale = 0.25
+_.bias_lr_mult = 2
+_.lr_mult = 4
 
 _.debug = True
 _.debug_sample_num = 1600

diff --git a/modules.py b/modules.py
@@ -51,7 +51,7 @@ def VGGBlock_ours(l):
 def Stage1Block(l, branch):
     assert branch in [1, 2]
 
-    ch_out = cfg.ch_heats if branch == 1 else cfg.ch_vectors
+    ch_out = cfg.ch_vectors if branch == 1 else cfg.ch_heats
 
     with tf.variable_scope('branch_%d' % branch):
         with argscope(Conv2D, W_init=tf.random_normal_initializer(stddev=0.01), nl=tf.nn.relu):
@@ -60,7 +60,6 @@ def Stage1Block(l, branch):
                  .Conv2D('conv2', 128, 3)
                  .Conv2D('conv3', 128, 3)
                  .Conv2D('conv4', 512, 1)
-                 # .Conv2D('conv5', ch_out, 1)())
                  .Conv2D('conv5', ch_out, 1, nl=tf.identity)())
 
     return l
@@ -70,7 +69,7 @@ def Stage1Block(l, branch):
 def StageTBlock(l, branch):
     assert branch in [1, 2]
 
-    ch_out = cfg.ch_heats if branch == 1 else cfg.ch_vectors
+    ch_out = cfg.ch_vectors if branch == 1 else cfg.ch_heats
 
     with tf.variable_scope('branch_%d' % branch):
         with argscope(Conv2D, W_init=tf.random_normal_initializer(stddev=0.01), nl=tf.nn.relu):
@@ -81,7 +80,6 @@ def StageTBlock(l, branch):
                  .Conv2D('conv4', 128, 7)
                  .Conv2D('conv5', 128, 7)
                  .Conv2D('conv6', 128, 1)
-                 # .Conv2D('conv7', ch_out, 1)())
                  .Conv2D('conv7', ch_out, 1, nl=tf.identity)())
 
     return l
diff --git a/train.py b/train.py
@@ -63,11 +63,11 @@ def _build_graph(self, inputs):
         if self.apply_mask:
             w1 = apply_mask(branch1, mask)
             w2 = apply_mask(branch2, mask)
-            heatmap_outputs.append(w1)
-            paf_outputs.append(w2)
+            paf_outputs.append(w1)
+            heatmap_outputs.append(w2)
         else:
-            heatmap_outputs.append(branch1)
-            paf_outputs.append(branch2)
+            paf_outputs.append(branch1)
+            heatmap_outputs.append(branch2)
 
         # Stage T
         for i in range(2, cfg.stages + 1):
@@ -77,11 +77,11 @@ def _build_graph(self, inputs):
             if self.apply_mask:
                 w1 = apply_mask(branch1, mask)
                 w2 = apply_mask(branch2, mask)
-                heatmap_outputs.append(w1)
-                paf_outputs.append(w2)
+                paf_outputs.append(w1)
+                heatmap_outputs.append(w2)
             else:
-                heatmap_outputs.append(branch1)
-                paf_outputs.append(branch2)
+                paf_outputs.append(branch1)
+                heatmap_outputs.append(branch2)
 
 
         # ========================== Cost Functions ==========================
@@ -132,7 +132,12 @@ def _get_optimizer(self):
         if cfg.backbone_grad_scale != 1.0:
             opt = tf.train.MomentumOptimizer(learning_rate=lr, momentum=cfg.momentum)
             gradprocs = [gradproc.ScaleGradient(
-                         [('conv[1-4]_[1-4]\/.*', cfg.backbone_grad_scale)])]
+                         [('conv.*/W', 1),
+                          ('conv.*/b', cfg.bias_lr_mult),
+                          ('stage_1.*/W', 1),
+                          ('stage_1.*/b', cfg.bias_lr_mult),
+                          ('stage_[2-6].*/W', cfg.lr_mult),
+                          ('stage_[2-6].*/b', cfg.lr_mult * cfg.bias_lr_mult)])]
             return optimizer.apply_grad_processors(opt, gradprocs)
         else:
             return tf.train.MomentumOptimizer(learning_rate=lr, momentum=cfg.momentum)