13c724a1b3dad1d5eddb51b851c39671982dbb69,sonnet/python/modules/conv.py,Conv2D,_build,#Conv2D#,335

Before Change

    self._input_channels = input_channels

    if not tf.float32.is_compatible_with(inputs.dtype):
      raise TypeError(
          "Input must have dtype tf.float32, but dtype was {}".format(
              inputs.dtype))

    weight_shape = (
        self._kernel_shape[0],
        self._kernel_shape[1],
        self._input_channels,
        self.output_channels)

    bias_shape = (self.output_channels,)

    if "w" not in self._initializers:
      self._initializers["w"] = create_weight_initializer(weight_shape[:3])

    if "b" not in self._initializers and self._use_bias:
      self._initializers["b"] = create_bias_initializer(bias_shape)

    self._w = tf.get_variable("w",
                              shape=weight_shape,
                              initializer=self._initializers["w"],
                              partitioner=self._partitioners.get("w", None),
                              regularizer=self._regularizers.get("w", None))

    w = self._w

    if self._mask is not None:
      mask = self._mask
      mask_shape = mask.shape.as_list()
      if len(mask_shape) == 2:
        if mask_shape != list(self._kernel_shape):
          raise base.IncompatibleShapeError(
              "Invalid mask shape: {}".format(tuple(mask_shape)))
        mask = tf.expand_dims(tf.expand_dims(mask, -1), -1)
      elif mask_shape != list(weight_shape):
        raise base.IncompatibleShapeError(
            "Invalid mask shape: {}".format(tuple(mask_shape)))
      w *= mask

    outputs = tf.nn.convolution(inputs, w, strides=self._stride,
                                padding=self._padding, dilation_rate=self._rate,
                                data_format=self._data_format)

    if self._use_bias:
      self._b = tf.get_variable("b",
                                shape=bias_shape,
                                initializer=self._initializers["b"],
                                partitioner=self._partitioners.get("b", None),
                                regularizer=self._regularizers.get("b", None))
      outputs = tf.nn.bias_add(outputs, self._b, data_format=self._data_format)

    return outputs

After Change

        self._input_channels,
        self.output_channels)

    bias_shape = (self.output_channels,)

    if "w" not in self._initializers:
      self._initializers["w"] = create_weight_initializer(weight_shape[:3],
                                                          dtype=inputs.dtype)

    if "b" not in self._initializers and self._use_bias:
      self._initializers["b"] = create_bias_initializer(bias_shape,
                                                        dtype=inputs.dtype)

    self._w = tf.get_variable("w",
                              shape=weight_shape,
                              dtype=inputs.dtype,
                              initializer=self._initializers["w"],
                              partitioner=self._partitioners.get("w", None),
                              regularizer=self._regularizers.get("w", None))

    w = self._w

    if self._mask is not None:
      mask = self._mask
      mask_shape = mask.shape.as_list()
      if len(mask_shape) == 2:
        if mask_shape != list(self._kernel_shape):
          raise base.IncompatibleShapeError(
              "Invalid mask shape: {}".format(tuple(mask_shape)))
        mask = tf.expand_dims(tf.expand_dims(mask, -1), -1)
      elif mask_shape != list(weight_shape):
        raise base.IncompatibleShapeError(
            "Invalid mask shape: {}".format(tuple(mask_shape)))
      w *= mask

    outputs = tf.nn.convolution(inputs, w, strides=self._stride,
                                padding=self._padding, dilation_rate=self._rate,
                                data_format=self._data_format)

    if self._use_bias:
      self._b = tf.get_variable("b",
                                shape=bias_shape,
                                dtype=inputs.dtype,
                                initializer=self._initializers["b"],
                                partitioner=self._partitioners.get("b", None),
                                regularizer=self._regularizers.get("b", None))
      outputs = tf.nn.bias_add(outputs, self._b, data_format=self._data_format)

    return outputs

Instances