Fix for certain effnet shapes with large model

pbaylies · Jun 16, 2019 · b56eb2d · b56eb2d
1 parent 5af40c0
commit b56eb2d
Showing 1 changed file with 5 additions and 11 deletions.
diff --git a/train_effnet.py b/train_effnet.py
@@ -181,19 +181,13 @@ def get_effnet_model(save_path, model_res=1024, image_size=256, depth=1, size=3,
             c = LocallyConnected1D(layer_l//3, 1, activation=activation)(c)
             x = Concatenate()([a,b,c])
         else:
-            a = LocallyConnected1D(layer_r//2, 1, activation=activation)(x)
-            b = LocallyConnected1D(layer_r//2, 1, activation=activation)(x)
-            c = LocallyConnected1D(layer_r//2, 1, activation=activation)(x)
-            d = LocallyConnected1D(layer_r//2, 1, activation=activation)(x)
+            a = LocallyConnected1D(layer_l, 1, activation=activation)(x)
+            b = LocallyConnected1D(layer_l, 1, activation=activation)(x)
             a = Permute((2, 1))(a)
             b = Permute((2, 1))(b)
-            c = Permute((2, 1))(c)
-            d = Permute((2, 1))(d)
-            a = LocallyConnected1D(layer_l//2, 1, activation=activation)(a)
-            b = LocallyConnected1D(layer_l//2, 1, activation=activation)(b)
-            c = LocallyConnected1D(layer_l//2, 1, activation=activation)(c)
-            d = LocallyConnected1D(layer_l//2, 1, activation=activation)(d)
-            x = Concatenate()([a,b,c,d])
+            a = LocallyConnected1D(layer_r//2, 1, activation=activation)(a)
+            b = LocallyConnected1D(layer_r//2, 1, activation=activation)(b)
+            x = Concatenate()([a,b])
         x = Add()([x, x_init])   # add skip connection
     x = Reshape((model_scale, 512))(x) # train against all dlatent values
     model = Model(inputs=inp,outputs=x)