Add pack layer tests, fix unpack_3d layer

This commit is contained in:
Piv
2021-07-19 20:37:54 +09:30
parent 38e7ad069e
commit 66cbc7faf6
2 changed files with 36 additions and 3 deletions

View File

@@ -68,12 +68,12 @@ def pack_3d(inputs, kernel_size, r=2, features_3d=8):
return packnet_conv2d(x, inputs.shape[3], kernel_size, 1)
def unpack_3d(inputs, out_channels, kernel_size, r=3, features_3d=8):
def unpack_3d(inputs, out_channels, kernel_size, r=2, features_3d=8):
x = packnet_conv2d(inputs, out_channels * (r ** 2) //
features_3d, kernel_size, 1)
x = tf.expand_dims(x, 4) # B x H/2 x W/2 x 4(out)/D x D
x = keras.layers.Conv3D(features_3d, kernel_size=3, padding='same')
b, c, d, h, w = x.shape
x = keras.layers.Conv3D(features_3d, kernel_size=3, padding='same')(x)
b, h, w, c, d = x.shape
x = tf.reshape(x, [b, h, w, c * d])
return nn.depth_to_space(x, r)