make network depth depend on img_size

sergio14890 · Nov 24, 2017 · c1d7626 · c1d7626
1 parent 5c16232
commit c1d7626
Showing 1 changed file with 10 additions and 4 deletions.
diff --git a/WRNs_imagenet.py b/WRNs_imagenet.py
@@ -194,10 +194,16 @@ def residual_block(l, increase_dim=False, first=False, filters=16):
         l = residual_block(l, filters=n_filters[2])
 
     # third stack of residual blocks
-
-    l = residual_block(l, increase_dim=True, filters=n_filters[3])
-    for _ in range(1, n):
-        l = residual_block(l, filters=n_filters[3])
+    if img_size >= 32:
+        l = residual_block(l, increase_dim=True, filters=n_filters[3])
+        for _ in range(1, n):
+            l = residual_block(l, filters=n_filters[3])
+
+    # fourth stack of residual blocks
+    if img_size >= 64:
+        l = residual_block(l, increase_dim=True, filters=n_filters[4])
+        for _ in range(1, n):
+            l = residual_block(l, filters=n_filters[4])
 
     bn_post_conv = BatchNormLayer(l)
     bn_post_relu = NonlinearityLayer(bn_post_conv, rectify)