Don't perform normalization on final layer

cswinter · Nov 9, 2019 · c03c2d0 · c03c2d0
1 parent 5f0f0c7
commit c03c2d0
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/hyper_params.py b/hyper_params.py
@@ -31,7 +31,7 @@ def __init__(self):
         self.obs_global_drones = 2  # Max number of (possibly hidden) drones observed by value function
         self.mconv_pooling = 'max'  # Pooling layer after mineral convolutions ('max', 'avg' or 'both')
         self.dconv_pooling = 'max'  # Pooling layer after drone convolutions ('max', 'avg' or 'both')
-        self.norm = 'batchnorm'     # Normalization layers ("none", "batchnorm", "layernorm")
+        self.norm = 'none'          # Normalization layers ("none", "batchnorm", "layernorm")
 
         # Eval
         self.eval_envs = 0

diff --git a/policy.py b/policy.py
@@ -94,7 +94,7 @@ def __init__(self,
                 kernel_size=1)
 
         layers = []
-        for _ in range(fc_layers - 1):
+        for i in range(fc_layers - 1):
             layers.append(
                 nn.Conv2d(
                     in_channels=nhidden,
@@ -103,7 +103,7 @@ def __init__(self,
                 )
             )
             layers.append(nn.ReLU())
-            if norm == 'none':
+            if norm == 'none' or i == fc_layers - 2:
                 pass
             elif norm == 'batchnorm':
                 layers.append(nn.BatchNorm2d(nhidden))