def model(self): d, s, m, r = self.model_params # Feature Extraction size = self.padding + 1 weights = tf.get_variable('w1', shape=[size, size, 1, d], initializer=tf.variance_scaling_initializer(0.1)) biases = tf.get_variable('b1', initializer=tf.zeros([d])) features = tf.nn.conv2d(self.images, weights, strides=[1,1,1,1], padding='VALID', data_format='NHWC') features = tf.nn.bias_add(features, biases, data_format='NHWC') # Shrinking if self.model_params[1] > 0: features = self.prelu(features, 1) weights = tf.get_variable('w2', shape=[1, 1, d, s], initializer=tf.variance_scaling_initializer(2)) biases = tf.get_variable('b2', initializer=tf.zeros([s])) features = tf.nn.conv2d(features, weights, strides=[1,1,1,1], padding='SAME', data_format='NHWC') features = tf.nn.bias_add(features, biases, data_format='NHWC') else: s = d conv = features # Mapping (# mapping layers = m) with tf.variable_scope("mapping_block") as scope: for ri in range(r): for i in range(3, m + 3): weights = tf.get_variable('w{}'.format(i), shape=[3, 3, s, s], initializer=tf.variance_scaling_initializer(2)) biases = tf.get_variable('b{}'.format(i), initializer=tf.zeros([s])) if i > 3: conv = self.prelu(conv, i) conv = tf.nn.conv2d(conv, weights, strides=[1,1,1,1], padding='SAME', data_format='NHWC') conv = tf.nn.bias_add(conv, biases, data_format='NHWC') if i == m + 2: conv = self.prelu(conv, m + 3) weights = tf.get_variable('w{}'.format(m + 3), shape=[1, 1, s, s], initializer=tf.variance_scaling_initializer(2)) biases = tf.get_variable('b{}'.format(m + 3), initializer=tf.zeros([s])) conv = tf.nn.conv2d(conv, weights, strides=[1,1,1,1], padding='SAME', data_format='NHWC') conv = tf.nn.bias_add(conv, biases, data_format='NHWC') conv = tf.add(conv, features) scope.reuse_variables() conv = self.prelu(conv, 2) # Expanding if self.model_params[1] > 0: expand_weights = tf.get_variable('w{}'.format(m + 4), shape=[1, 1, s, d], initializer=tf.variance_scaling_initializer(2)) expand_biases = tf.get_variable('b{}'.format(m + 4), initializer=tf.zeros([d])) conv = tf.nn.conv2d(conv, expand_weights, strides=[1,1,1,1], padding='SAME', data_format='NHWC') conv = tf.nn.bias_add(conv, expand_biases, data_format='NHWC') conv = self.prelu(conv, m + 4) # Sub-pixel convolution size = self.radius * 2 + 1 deconv_weights = tf.get_variable('deconv_w', shape=[size, size, d, self.scale**2], initializer=tf.variance_scaling_initializer(0.01)) deconv_biases = tf.get_variable('deconv_b', initializer=tf.zeros([self.scale**2])) deconv = tf.nn.conv2d(conv, deconv_weights, strides=[1,1,1,1], padding='SAME', data_format='NHWC') deconv = tf.nn.bias_add(deconv, deconv_biases, data_format='NHWC') if self.scale > 1: deconv = tf.depth_to_space(deconv, self.scale, name='pixel_shuffle', data_format='NHWC') return deconv
def sk_conv2d(inputs, filters, strides, sk_ratio, min_dim=32, is_training=True, data_format='channels_last'): """Selective kernel convolutional layer (https://arxiv.org/abs/1903.06586).""" channel_axis = 1 if data_format == 'channels_first' else 3 pooling_axes = [2, 3] if data_format == 'channels_first' else [1, 2] # Two stream convs (using split and both are 3x3). inputs = conv2d_fixed_padding( inputs=inputs, filters=2 * filters, kernel_size=3, strides=strides, data_format=data_format) inputs = batch_norm_relu(inputs, is_training, data_format=data_format) inputs = tf.stack(tf.split(inputs, num_or_size_splits=2, axis=channel_axis)) # Mixing weights for two streams. mid_dim = max(int(filters * sk_ratio), min_dim) global_features = tf.reduce_mean( tf.reduce_sum(inputs, axis=0), pooling_axes, keepdims=True) global_features = tf.layers.conv2d( inputs=global_features, filters=mid_dim, kernel_size=1, strides=1, kernel_initializer=tf.variance_scaling_initializer(), use_bias=False, data_format=data_format) global_features = batch_norm_relu( global_features, is_training, data_format=data_format) mixing = tf.layers.conv2d( inputs=global_features, filters=2 * filters, kernel_size=1, strides=1, kernel_initializer=tf.variance_scaling_initializer(), use_bias=False, data_format=data_format) mixing = tf.stack(tf.split(mixing, num_or_size_splits=2, axis=channel_axis)) mixing = tf.nn.softmax(mixing, axis=0) return tf.reduce_sum(inputs * mixing, axis=0)
def se_layer(inputs, filters, se_ratio, data_format='channels_last'): """Squeeze and Excitation layer (https://arxiv.org/abs/1709.01507).""" if se_ratio <= 0: return inputs se_reduce = tf.layers.Conv2D( max(1, int(filters * se_ratio)), kernel_size=[1, 1], strides=[1, 1], kernel_initializer=tf.variance_scaling_initializer(), padding='same', data_format=data_format, use_bias=True) se_expand = tf.layers.Conv2D( inputs.shape[-1], kernel_size=[1, 1], strides=[1, 1], kernel_initializer=tf.variance_scaling_initializer(), padding='same', data_format=data_format, use_bias=True) spatial_dims = [2, 3] if data_format == 'channels_first' else [1, 2] se_tensor = tf.reduce_mean( inputs, spatial_dims, keepdims=True) se_tensor = se_expand(tf.nn.relu(se_reduce(se_tensor))) return tf.sigmoid(se_tensor) * inputs
def get_variable_initializer(hparams): """Get variable initializer from hparams.""" if not hparams.initializer: return None mlperf_log.transformer_print(key=mlperf_log.MODEL_HP_INITIALIZER_GAIN, value=hparams.initializer_gain, hparams=hparams) if not tf.executing_eagerly(): tf.logging.info("Using variable initializer: %s", hparams.initializer) if hparams.initializer == "orthogonal": return tf.orthogonal_initializer(gain=hparams.initializer_gain) elif hparams.initializer == "uniform": max_val = 0.1 * hparams.initializer_gain return tf.random_uniform_initializer(-max_val, max_val) elif hparams.initializer == "normal_unit_scaling": return tf.variance_scaling_initializer( hparams.initializer_gain, mode="fan_avg", distribution="normal") elif hparams.initializer == "uniform_unit_scaling": return tf.variance_scaling_initializer( hparams.initializer_gain, mode="fan_avg", distribution="uniform") elif hparams.initializer == "xavier": return tf.initializers.glorot_uniform() else: raise ValueError("Unrecognized initializer: %s" % hparams.initializer)
def __init__(self, num_inputs, num_outputs, data_dim, activation=tf.nn.relu, name="resmade_block"): self.num_inputs = num_inputs self.num_outputs = num_outputs self.activation = activation with tf.variable_scope(name): self.layer1 = MaskedDenseLayer( num_inputs, num_outputs, data_dim, activation=activation, mask_type="hidden", weight_initializer=tf.variance_scaling_initializer( scale=2.0, distribution="normal"), name="layer1") self.layer2 = MaskedDenseLayer( num_inputs, num_outputs, data_dim, activation=None, mask_type="hidden", weight_initializer=tf.variance_scaling_initializer( scale=0.1, distribution="normal"), name="layer2")
def build_network(self): with tf.compat.v1.variable_scope(self.name): self.input = tf.placeholder(tf.float32, shape=[None, *self.input_dims], name='inputs') self.actions = tf.placeholder(tf.float32, shape=[None, self.n_actions], name='action_taken') self.q_target = tf.placeholder(tf.float32, shape=[None, self.n_actions], name='q_value') conv1 = tf.layers.conv2d( inputs=self.input, filters=32, kernel_size=(8, 8), strides=4, name='conv1', kernel_initializer=tf.variance_scaling_initializer(scale=2)) conv1_activated = tf.nn.relu(conv1) conv2 = tf.layers.conv2d( inputs=conv1_activated, filters=64, kernel_size=(4, 4), strides=2, name='conv2', kernel_initializer=tf.variance_scaling_initializer(scale=2)) conv2_activated = tf.nn.relu(conv2) conv3 = tf.layers.conv2d( inputs=conv2_activated, filters=128, kernel_size=(3, 3), strides=1, name='conv3', kernel_initializer=tf.variance_scaling_initializer(scale=2)) conv3_activated = tf.nn.relu(conv3) flat = tf.layers.flatten(conv3_activated) dense1 = tf.layers.dense( flat, units=self.fc1_dims, activation=tf.nn.relu, kernel_initializer=tf.variance_scaling_initializer(scale=2)) self.Q_values = tf.layers.dense( dense1, units=self.n_actions, kernel_initializer=tf.variance_scaling_initializer(scale=2)) #self.q = tf.reduce_sum(tf.multiply(self.Q_values, self.actions)) self.loss = tf.reduce_mean(tf.square(self.Q_values - self.q_target)) self.train_op = tf.train.AdamOptimizer(self.lr).minimize(self.loss)
def model(self): d = self.model_params m = len(d) + 2 # Feature Extraction size = self.padding + 1 weights = tf.get_variable( 'w1', shape=[size, size, 1, d[0]], initializer=tf.variance_scaling_initializer(0.1)) biases = tf.get_variable('b1', initializer=tf.zeros([d[0]])) conv = tf.nn.conv2d(self.images, weights, strides=[1, 1, 1, 1], padding='VALID', data_format='NHWC') conv = tf.nn.bias_add(conv, biases, data_format='NHWC') conv = self.prelu(conv, 1) # Mapping (# mapping layers = m) for i in range(3, m): weights = tf.get_variable( 'w{}'.format(i), shape=[3, 3, d[i - 3], d[i - 2]], initializer=tf.variance_scaling_initializer(2)) biases = tf.get_variable('b{}'.format(i), initializer=tf.zeros([d[i - 2]])) conv = tf.nn.conv2d(conv, weights, strides=[1, 1, 1, 1], padding='SAME', data_format='NHWC') conv = tf.nn.bias_add(conv, biases, data_format='NHWC') conv = self.prelu(conv, i) # Sub-pixel convolution size = self.radius * 2 + 1 deconv_weights = tf.get_variable( 'deconv_w', shape=[size, size, d[-1], self.scale**2], initializer=tf.variance_scaling_initializer(0.01)) deconv_biases = tf.get_variable('deconv_b', initializer=tf.zeros([self.scale**2])) deconv = tf.nn.conv2d(conv, deconv_weights, strides=[1, 1, 1, 1], padding='SAME', data_format='NHWC') deconv = tf.nn.bias_add(deconv, deconv_biases, data_format='NHWC') deconv = tf.depth_to_space(deconv, self.scale, name='pixel_shuffle', data_format='NHWC') return deconv
def __call__(self, input, **kwargs): """Call separable_conv2d function.""" with tf.variable_scope(self._scope_name, reuse=tf.AUTO_REUSE): return tf.keras.layers.SeparableConv2D(filters=self.out_channels, kernel_size=self.kernel_size, strides=self.stride, data_format=self.data_format, dilation_rate=self.dilation, depthwise_initializer=tf.variance_scaling_initializer(), pointwise_initializer=tf.variance_scaling_initializer(), padding='SAME', use_bias=self.bias, name='SeparableConv2d', reuse=self.reuse)(inputs=input)
def get_initializer(initializer, dtype): if initializer == 'zeros': return tf.zeros_initializer(dtype=dtype) elif initializer == 'ones': return tf.ones_initializer(dtype=dtype) elif initializer == 'vs': return tf.variance_scaling_initializer(dtype=dtype) elif initializer == 'xavier': return tf.glorot_normal_initializer(dtype=dtype) elif initializer == 'he': return tf.variance_scaling_initializer(dtype=dtype) else: raise NotImplementedError
def call(self, input, **kwargs): """Call separable_conv2d function.""" model = tf.keras.layers.SeparableConv2D(filters=self.out_channels, kernel_size=self.kernel_size, strides=self.stride, data_format=self.data_format, dilation_rate=self.dilation, depthwise_initializer=tf.variance_scaling_initializer(), pointwise_initializer=tf.variance_scaling_initializer(), padding='SAME', use_bias=self.bias, name=self.name, reuse=self.reuse, trainable=self._trainable) return model(inputs=input)
def squeeze_excitation(inputs, in_filters, se_ratio, name=None, expand_ratio=1, data_format='channels_last'): """Squeeze and excitation implementation. Args: inputs: `Tensor` of size `[batch, channels, height_in, width_in]`. in_filters: `int` number of input filteres before expansion. se_ratio: `float` a se ratio between 0 and 1 for squeeze and excitation. expand_ratio: `int` expansion ratio for the block. data_format: An optional string from: "channels_last", "channels_first". Defaults to "channels_last". Returns: A `Tensor` of shape `[batch, filters, height_out, width_out]`. """ num_reduced_filters = max(1, int(in_filters * se_ratio)) se_reduce = tf.layers.Conv2D( num_reduced_filters, kernel_size=[1, 1], strides=[1, 1], kernel_initializer=tf.variance_scaling_initializer(), padding='same', data_format=data_format, name="se0_{}".format(name), use_bias=True) se_expand = tf.layers.Conv2D( in_filters * expand_ratio, kernel_size=[1, 1], strides=[1, 1], kernel_initializer=tf.variance_scaling_initializer(), padding='same', data_format=data_format, name="se1_{}".format(name), use_bias=True) # Process input if data_format == 'channels_first': spatial_dims = [2, 3] else: spatial_dims = [1, 2] se_tensor = tf.reduce_mean(inputs, spatial_dims, keepdims=True) se_tensor = se_expand(tf.nn.relu(se_reduce(se_tensor))) return tf.sigmoid(se_tensor) * inputs
def __init__(self, scope=None): self.weights={} self.scope=scope self.kernel_initializer=tf.variance_scaling_initializer() self.build_CNN_params() print('Initialize weights {}'.format(self.scope))
def get_initializer(initializer, dtype): if initializer == 'zeros': return tf.zeros_initializer(dtype=dtype) elif initializer == 'vs': return tf.variance_scaling_initializer(dtype=dtype) #初始化 else: raise NotImplementedError
def conv_linear_map(inputs, nin, nout, bias_start, prefix): """Convolutional liner map. Maps 3D tensor by last dimension. Args: inputs: Inputs that should be shuffled nin: Input feature map count nout: Output feature map count bias_start: Bias start value prefix: Name prefix Returns: tf.Tensor: Inputs with applied convolution """ with tf.variable_scope(prefix): inp_shape = tf.shape(inputs) initializer = tf.variance_scaling_initializer( scale=1.0, mode="fan_avg", distribution="uniform") kernel = tf.get_variable("CvK", [nin, nout], initializer=initializer) bias_term = tf.get_variable( "CvB", [nout], initializer=tf.constant_initializer(0.0)) mul_shape = [inp_shape[0] * inp_shape[1], nin] res = tf.matmul(tf.reshape(inputs, mul_shape), kernel) res = tf.reshape(res, [inp_shape[0], inp_shape[1], nout]) return res + bias_start + bias_term
def create_conv_layer(filter_size, stride, input_size, output_size, input): """ :param filter_size: size of the filter :param stride: stride :param input_size: input size :param output_size: number of neurons :param input: input of the layer :return A_conv: ouput """ with tf.variable_scope(None, default_name="conv"): # Initialize the weights with tf.variance_scaling_initializer W_init = tf.variance_scaling_initializer() W_conv = tf.Variable( W_init([filter_size, filter_size, input_size, output_size])) # Initialize the biases with tf.zeros initializer b_init = tf.zeros_initializer() b_conv = tf.Variable(b_init(output_size)) A_conv = tf.nn.conv2d( input, W_conv, strides=[1, stride, stride, 1 ], padding='SAME') + b_conv # Apply activation function A_conv = tf.nn.relu(A_conv) return A_conv
def prn(x, is_training): """ Arguments: x: a float tensor with shape [b, h, w, c]. is_training: a boolean. Returns: a float tensor with shape [b, h, w, c]. """ with tf.variable_scope('PRN'): b = tf.shape(x)[0] _, h, w, c = x.shape.as_list() # must be static x = tf.reshape(x, [b, h * w * c]) # flatten with slim.arg_scope( [slim.fully_connected], weights_initializer=tf.variance_scaling_initializer()): y = slim.fully_connected(x, 1024, activation_fn=tf.nn.relu, scope='fc1') # y = slim.dropout(y, keep_prob=0.5, is_training=is_training) y = slim.fully_connected(y, h * w * c, activation_fn=tf.nn.relu, scope='fc2') x += y return tf.reshape(x, [b, h, w, c])
def conv2d_fixed_padding(inputs, filters, kernel_size, strides, data_format='channels_first'): """Strided 2-D convolution with explicit padding. The padding is consistent and is based only on `kernel_size`, not on the dimensions of `inputs` (as opposed to using `tf.layers.conv2d` alone). Args: inputs: `Tensor` of size `[batch, channels, height_in, width_in]`. filters: `int` number of filters in the convolution. kernel_size: `int` size of the kernel to be used in the convolution. strides: `int` strides of the convolution. data_format: `str` either "channels_first" for `[batch, channels, height, width]` or "channels_last for `[batch, height, width, channels]`. Returns: A `Tensor` of shape `[batch, filters, height_out, width_out]`. """ if strides > 1: inputs = fixed_padding(inputs, kernel_size, data_format=data_format) return tf.layers.conv2d( inputs=inputs, filters=filters, kernel_size=kernel_size, strides=strides, padding=('SAME' if strides == 1 else 'VALID'), use_bias=False, kernel_initializer=tf.variance_scaling_initializer(), data_format=data_format)
def conv_bn_relu(inputs, conv_size, conv_filters, is_training, data_format): """Convolution followed by batch norm and ReLU.""" if data_format == 'channels_last': axis = 3 elif data_format == 'channels_first': axis = 1 else: raise ValueError('invalid data_format') net = tf.layers.conv2d( inputs=inputs, filters=conv_filters, kernel_size=conv_size, strides=(1, 1), use_bias=False, kernel_initializer=tf.variance_scaling_initializer(), padding='same', data_format=data_format) net = tf.layers.batch_normalization(inputs=net, axis=axis, momentum=BN_MOMENTUM, epsilon=BN_EPSILON, training=is_training) net = tf.nn.relu(net) return net
def se_block(input_feature, name, ratio=8): kernel_initializer = tf.variance_scaling_initializer() bias_initializer = tf.constant_initializer(value=0.0) with tf.variable_scope(name): channel = input_feature.get_shape()[-1] # Global average pooling squeeze = tf.reduce_mean(input_feature, axis=[1,2], keepdims=True) assert squeeze.get_shape()[1:] == (1,1,channel) excitation = tf.layers.dense(inputs=squeeze, units=channel//ratio, activation=tf.nn.relu, kernel_initializer=kernel_initializer, bias_initializer=bias_initializer, name='bottleneck_fc', trainable=False) assert excitation.get_shape()[1:] == (1,1,channel//ratio) excitation = tf.layers.dense(inputs=excitation, units=channel, activation=tf.nn.sigmoid, kernel_initializer=kernel_initializer, bias_initializer=bias_initializer, name='recover_fc', trainable=False) assert excitation.get_shape()[1:] == (1,1,channel) scale = input_feature * excitation return scale
def build(self, input_shape): """Initialize layer weights and sublayers. Args: input_shape: shape of inputs """ in_units = input_shape[-1] middle_units = in_units * 4 out_units = in_units * 2 init = tf.variance_scaling_initializer(scale=1.0, mode="fan_avg", distribution="uniform") self.first_linear = tf.keras.layers.Dense(middle_units, use_bias=False, kernel_initializer=init, name=self.prefix + "/cand1") self.second_linear = tf.keras.layers.Dense(out_units, kernel_initializer=init, name=self.prefix + "/cand2") self.layer_norm = LayerNormalization() init = tf.constant_initializer(self.init_value) self.residual_scale = self.add_weight(self.prefix + "/residual", [out_units], initializer=init) super(RSU, self).build(input_shape)
def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=0, bias=True, groups=1, dilation=1, separable=False, depthwise=False, padding_mode='same', bn=False): super(Conv2d, self).__init__() self.in_channels = in_channels self.out_channels = out_channels self.kernel_size = kernel_size self.stride = stride self.padding = padding self.bias = bias self.groups = groups self.dilation = dilation self.kernel_initial = tf.variance_scaling_initializer() self.bias_initial = tf.zeros_initializer() self._initializer = None self.reuse = None self.separable = separable self.depthwise = depthwise self.padding_mode = padding if isinstance( padding, str) and stride != 2 else padding_mode self.bn = bn
def initialize_parameters(): tf.set_random_seed(1) W1 = tf.get_variable("W1", [25, 12288], initializer=tf.variance_scaling_initializer(seed=1)) b1 = tf.get_variable("b1", [25, 1], initializer=tf.zeros_initializer()) W2 = tf.get_variable("W2", [12, 25], initializer=tf.variance_scaling_initializer(seed=1)) b2 = tf.get_variable("b2", [12, 1], initializer=tf.zeros_initializer()) W3 = tf.get_variable("W3", [6, 12], initializer=tf.variance_scaling_initializer(seed=1)) b3 = tf.get_variable("b3", [6, 1], initializer=tf.zeros_initializer()) parameters = {"W1": W1, "b1": b1, "W2": W2, "b2": b2, "W3": W3, "b3": b3} return parameters
def _network_template(self, obs, num_layers, hidden_units): """PixelCNN network architecture.""" with slim.arg_scope( [slim.conv2d, masked_conv2d], weights_initializer=tf.variance_scaling_initializer( distribution='uniform'), biases_initializer=tf.constant_initializer(0.0)): net = masked_conv2d(obs, hidden_units, [7, 7], mask_type='A', activation_fn=None, scope='masked_conv_1') embedding = slim.model_variable( 'embedding', shape=(1, ) + self.resize_shape + (4, ), initializer=tf.variance_scaling_initializer( distribution='uniform')) for i in range(1, num_layers + 1): net2 = gating_layer(net, embedding, hidden_units, 'gating_{}'.format(i)) net += masked_conv2d(net2, hidden_units, [1, 1], mask_type='B', activation_fn=None, scope='masked_conv_{}'.format(i + 1)) net += slim.conv2d(embedding, hidden_units, [1, 1], activation_fn=None) net = tf.nn.relu(net) net = masked_conv2d(net, 64, [1, 1], scope='1x1_conv_out', mask_type='B', activation_fn=tf.nn.relu) logits = masked_conv2d(net, self.quantization_factor, [1, 1], scope='logits', mask_type='B', activation_fn=None) loss = tf.losses.sparse_softmax_cross_entropy( labels=tf.cast(obs, tf.int32), logits=logits, reduction=tf.losses.Reduction.MEAN) return collections.namedtuple('PixelCNN_network', ['logits', 'loss'])(logits, loss)
def conv0_space_to_depth(inputs, filters, kernel_size, strides, data_format='channels_last', space_to_depth_block_size=2): """Uses space-to-depth convolution for conv0. This function replaces the first convolution (conv0) in ResNet with space-to-depth transformation. It creates a convolution kernel, whose dimension and name are the same as those of conv0. The `inputs` is an image tensor that already has the space-to-depth transform. Args: inputs: `Tensor` of size `[batch, height_in, width_in, channels]`. filters: An `int` number of filters in the convolution. kernel_size: An `int` size of the kernel to be used in the convolution. strides: A `int` strides of the convolution. data_format: A `str` either "channels_first" for `[batch, channels, height, width]` or "channels_last for `[batch, height, width, channels]`. space_to_depth_block_size: An `int` indicates the block size of space-to-depth convolution for conv0. Specific to ResNet, this currently supports only block_size=2. Returns: A `Tensor` with the same type as `inputs`. Raises: ValueError if `space_to_depth_block_size` is not 2. """ if space_to_depth_block_size != 2: raise ValueError('Space-to-depth does not support block_size (%d).' % space_to_depth_block_size) conv0 = tf.layers.Conv2D( filters=filters, kernel_size=kernel_size, strides=strides, padding=('SAME' if strides == 1 else 'VALID'), data_format=data_format, use_bias=False, kernel_initializer=tf.variance_scaling_initializer()) # Use the image size without space-to-depth transform as the input of conv0. # This allows the kernel size to be the same as the original conv0 such that # the model is able to load the pre-trained ResNet checkpoint. batch_size, h, w, c = inputs.get_shape().as_list() conv0.build([batch_size, h * space_to_depth_block_size, w * space_to_depth_block_size, c // (space_to_depth_block_size ** 2)]) kernel = conv0.weights[0] kernel = transform_space_to_depth_kernel( kernel, inputs.dtype, block_size=space_to_depth_block_size) inputs = spatial_transform.space_to_depth_fixed_padding( inputs, kernel_size, data_format, space_to_depth_block_size) return tf.nn.conv2d( input=inputs, filter=kernel, strides=[1, 1, 1, 1], padding='VALID', data_format='NHWC' if data_format == 'channels_last' else 'NCHW', name='conv2d/Conv2D')
def conv2d_fixed_padding(inputs, filters, kernel_size, strides, pruning_method='baseline', data_format='channels_first', weight_decay=0., name=None): """Strided 2-D convolution with explicit padding. The padding is consistent and is based only on `kernel_size`, not on the dimensions of `inputs` (as opposed to using `tf.layers.conv2d` alone). Args: inputs: Input tensor, float32 or bfloat16 of size [batch, channels, height, width]. filters: Int specifying number of filters for the first two convolutions. kernel_size: Int designating size of kernel to be used in the convolution. strides: Int specifying the stride. If stride >1, the input is downsampled. pruning_method: String that specifies the pruning method used to identify which weights to remove. data_format: String that specifies either "channels_first" for [batch, channels, height,width] or "channels_last" for [batch, height, width, channels]. weight_decay: Weight for the l2 regularization loss. name: String that specifies name for model layer. Returns: The output activation tensor of size [batch, filters, height_out, width_out] Raises: ValueError: If the data_format provided is not a valid string. """ if strides > 1: inputs = resnet_model.fixed_padding(inputs, kernel_size, data_format=data_format) padding = 'VALID' else: padding = 'SAME' kernel_initializer = tf.variance_scaling_initializer() kernel_regularizer = contrib_layers.l2_regularizer(weight_decay) return sparse_conv2d(x=inputs, units=filters, activation=None, kernel_size=[kernel_size, kernel_size], use_bias=False, kernel_initializer=kernel_initializer, kernel_regularizer=kernel_regularizer, bias_initializer=None, biases_regularizer=None, sparsity_technique=pruning_method, normalizer_fn=None, strides=[strides, strides], padding=padding, data_format=data_format, name=name)
def conv0_space_to_depth(inputs, data_format='channels_last'): """Strided 2-D convolution with explicit padding. The padding is consistent and is based only on `kernel_size`, not on the dimensions of `inputs` (as opposed to using `tf.layers.conv2d` alone). Args: inputs: `Tensor` of size `[batch, height_in, width_in, channels]`. data_format: `str` either "channels_first" for `[batch, channels, height, width]` or "channels_last for `[batch, height, width, channels]`. Returns: A `Tensor` with the same type as `inputs`. """ # Create the conv0 kernel w.r.t. the original image size. (no space-to-depth). filters = 64 kernel_size = 7 space_to_depth_block_size = ssd_constants.SPACE_TO_DEPTH_BLOCK_SIZE strides = 2 conv0 = tf.compat.v1.layers.Conv2D( filters=filters, kernel_size=kernel_size, strides=2, padding=('SAME' if strides == 1 else 'VALID'), use_bias=False, kernel_initializer=tf.variance_scaling_initializer(), data_format=data_format) # Use the image size without space-to-depth transform as the input of conv0. batch_size, h, w, channel = inputs.get_shape().as_list() conv0.build([ batch_size, h * space_to_depth_block_size, w * space_to_depth_block_size, channel // (space_to_depth_block_size**2) ]) kernel = conv0.weights[0] # [7, 7, 3, 64] --> [8, 8, 3, 64] kernel = tf.pad(kernel, paddings=tf.constant([[1, 0], [1, 0], [0, 0], [0, 0]]), mode='CONSTANT', constant_values=0.) # Transform kernel follows the space-to-depth logic: http://shortn/_9YvHW96xPJ kernel = tf.reshape(kernel, [ 4, space_to_depth_block_size, 4, space_to_depth_block_size, 3, filters ]) kernel = tf.transpose(kernel, [0, 2, 1, 3, 4, 5]) kernel = tf.reshape(kernel, [4, 4, int(channel), filters]) kernel = tf.cast(kernel, inputs.dtype) inputs = space_to_depth_fixed_padding(inputs, kernel_size, data_format, space_to_depth_block_size) return tf.nn.conv2d( input=inputs, filter=kernel, strides=[1, 1, 1, 1], padding='VALID', data_format='NHWC' if data_format == 'channels_last' else 'NCHW', name='conv2d/Conv2D')
def __init__(self, sparsity, seed=None, dtype=tf.float32): if sparsity < 0. or sparsity > 1.: raise ValueError('sparsity must be in the range [0., 1.].') self.kernel_initializer = tf.variance_scaling_initializer(seed=seed, dtype=dtype) self.seed = seed self.dtype = dtype self.sparsity = float(sparsity)
def conv2d_fixed_padding(inputs, filters, kernel_size, strides, data_format): if strides > 1: inputs = fixed_padding(inputs, kernel_size, data_format) return tf.layers.conv2d( inputs=inputs, filters=filters, kernel_size=kernel_size, strides=strides, padding=('SAME' if strides == 1 else 'VALID'), use_bias=False, kernel_initializer=tf.variance_scaling_initializer(), data_format=data_format)
def conv2d_space_to_depth(inputs, filters, kernel_size, strides, block_size=2): """Strided 2-D convolution with explicit padding. The padding is consistent and is based only on `kernel_size`, not on the dimensions of `inputs` (as opposed to using `tf.layers.conv2d` alone). Args: inputs: `Tensor` of size `[batch, height_in, width_in, channels]`. filters: `int` number of filters in the convolution. kernel_size: `int` size of the kernel to be used in the convolution. strides: `int` strides of the convolution. block_size: space-to-depth padding block-size Returns: A `Tensor` of shape `[batch, filters, height_out, width_out]`. """ def _round_up(kernel_size, block_size): remainder = kernel_size % block_size if remainder == 0: return block_size else: return kernel_size + block_size - remainder padded_k = _round_up(kernel_size, block_size) _, _, _, in_f = inputs.get_shape().as_list() # Create the padded kernel kernel = tf.get_variable( 'kernel_conv2d_opt', shape=[padded_k, padded_k, in_f // 4, filters], initializer=tf.variance_scaling_initializer(), trainable=True, dtype=tf.float32) # Zero padded region mx = tf.constant([1, 1, 1, 1, 1, 1, 1, 0], dtype=tf.float32) kernel = kernel * tf.reshape(mx, [1, padded_k, 1, 1]) * tf.reshape( mx, [padded_k, 1, 1, 1]) # Transpose to enable space-to-depth optimization kernel = tf.reshape(kernel, [ padded_k // 2, block_size, padded_k // 2, block_size, in_f // 4, filters ]) kernel = tf.transpose(kernel, [0, 2, 1, 3, 4, 5]) kernel = tf.reshape(kernel, [padded_k // 2, padded_k // 2, in_f, filters]) kernel = tf.cast(kernel, inputs.dtype) # Add a convolution operation using this kernel s = strides // 2 return tf.nn.conv2d( input=inputs, filter=kernel, strides=[s, s, s, s], padding='SAME', data_format='NHWC')
def default_mlp(hidden_sizes, activate_final=False, init_std=2., **kwargs): """Standard batch-applied MLP for transformer modules.""" init = { 'w': tf.variance_scaling_initializer(init_std, distribution='normal') } mlp = snt_mlp.MLP(hidden_sizes, activate_final=activate_final, use_dropout=True, initializers=init, **kwargs) return basic.BatchApply(mlp)