def test_pass_through(self, width, height, filters, mod): x = tf.keras.Input(shape=(width, height, filters)) test_layer = nn_blocks.CSPRoute(filters=filters, filter_scale=mod) test_layer2 = nn_blocks.CSPConnect(filters=filters, filter_scale=mod) outx, px = test_layer(x) outx = test_layer2([outx, px]) print(outx) print(outx.shape.as_list()) self.assertAllEqual( outx.shape.as_list(), [None, np.ceil(width // 2), np.ceil(height // 2), (filters)])
def _csp_stack(self, inputs, config, name): if config.bottleneck: csp_filter_scale = 1 residual_filter_scale = 2 scale_filters = 1 else: csp_filter_scale = 2 residual_filter_scale = 1 scale_filters = 2 self._default_dict['activation'] = self._get_activation( config.activation) self._default_dict['name'] = f'{name}_csp_down' if self._dilate: self._default_dict['dilation_rate'] = config.dilation_rate degrid = int( tf.math.log(float(config.dilation_rate)) / tf.math.log(2.)) else: self._default_dict['dilation_rate'] = 1 degrid = 0 # swap/add dialation x, x_route = nn_blocks.CSPRoute(filters=config.filters, filter_scale=csp_filter_scale, downsample=True, **self._default_dict)(inputs) dilated_reps = config.repetitions - degrid for i in range(dilated_reps): self._default_dict['name'] = f'{name}_{i}' x = nn_blocks.DarkResidual(filters=config.filters // scale_filters, filter_scale=residual_filter_scale, **self._default_dict)(x) for i in range(dilated_reps, config.repetitions): self._default_dict['dilation_rate'] = max( 1, self._default_dict['dilation_rate'] // 2) self._default_dict[ 'name'] = f"{name}_{i}_degridded_{self._default_dict['dilation_rate']}" x = nn_blocks.DarkResidual(filters=config.filters // scale_filters, filter_scale=residual_filter_scale, **self._default_dict)(x) self._default_dict['name'] = f'{name}_csp_connect' output = nn_blocks.CSPConnect(filters=config.filters, filter_scale=csp_filter_scale, **self._default_dict)([x, x_route]) self._default_dict['activation'] = self._activation self._default_dict['name'] = None return output
def test_gradient_pass_though(self, filters, width, height, mod): loss = tf.keras.losses.MeanSquaredError() optimizer = tf.keras.optimizers.SGD() test_layer = nn_blocks.CSPRoute(filters, filter_scale=mod) path_layer = nn_blocks.CSPConnect(filters, filter_scale=mod) init = tf.random_normal_initializer() x = tf.Variable(initial_value=init(shape=(1, width, height, filters), dtype=tf.float32)) y = tf.Variable(initial_value=init(shape=(1, int(np.ceil(width // 2)), int(np.ceil(height // 2)), filters), dtype=tf.float32)) with tf.GradientTape() as tape: x_hat, x_prev = test_layer(x) x_hat = path_layer([x_hat, x_prev]) grad_loss = loss(x_hat, y) grad = tape.gradient(grad_loss, test_layer.trainable_variables) optimizer.apply_gradients(zip(grad, test_layer.trainable_variables)) self.assertNotIn(None, grad)