def compute_self_att_v2(
    end_points,
    logits,
    num_class=21,
    attention_dim=128,
    strides=(15, 16),
    is_training=True,
    linformer=True,
    valid_mask=None,
    factor=8,
    downsample_type='nearest',
    net='xception_65'):
  """Compute self-attention for segmentation head.

  Args:
    end_points: Network end_points (dict).
    logits: The input seed for refinement. Used as ``value'' in self-attention.
    Can be either logits, probability, or score map.
    num_class: Number of classes including background
    attention_dim: Embedding space dimension for key and query used in the
    self-attention module
    strides: Use feature maps from which stride to compute pixel similarity
    is_training: Indicate training or inference mode
    linformer: Adopt the idea from https://arxiv.org/abs/2006.04768 to reduce
    memory usage in self-attention computation. But instead of learning the
    downsample function, we use deterministic image downsample functions
    valid_mask: To identity valid region of the input. It is used to avoid
    attending to padding regions
    factor: Downsample factor used in linformer mode
    downsample_type: Use which downsample method to reduce the memory usage. Can
    be either 'nearest' or 'bilinear'. Default: 'nearest'
    net: Specify which network is used

  Returns:
    A list of computed Grad-CAMs or refined ones.
  """
  # Sanity check: Make sure strides are sorted
  strides = sorted(list(strides))[::-1]
  conv_layer_list = []
  for stride in strides:
    conv_layer = end_points[net_to_stride_to_endpoints_name[net][stride]]
    conv_layer_list.append(conv_layer)

  # Resize to seed resolution first
  h, w = preprocess_utils.resolve_shape(logits, 4)[1:3]
  conv_layer_list = [
      tf.compat.v1.image.resize_bilinear(
          conv, (h, w), align_corners=True)
      for conv in conv_layer_list
  ]
  conv_layer_merged = tf.concat(conv_layer_list, axis=-1)
  conv_layer_merged = tf.stop_gradient(conv_layer_merged)
  score = tf.stop_gradient(logits)
  # This tells us what input it is (decoder logits or Grad-CAM)
  value_dim = tf.shape(score)[-1]

  # Only valid when we use Linformer style to reduce size for key and value
  if downsample_type == 'bilinear':
    resize_fn = tf.compat.v1.image.resize_bilinear
  else:
    resize_fn = tf.compat.v1.image.resize_nearest_neighbor

  scope = 'hyper_column'
  with tf.variable_scope(scope):
    with slim.arg_scope([slim.conv2d],
                        activation_fn=None,
                        normalizer_fn=None,
                        biases_initializer=None,
                        reuse=tf.AUTO_REUSE):
      k = slim.conv2d(
          conv_layer_merged, attention_dim, [1, 1], scope='key')
      q = slim.conv2d(
          conv_layer_merged, attention_dim, [1, 1], scope='query')
  q = tf.reshape(q, [-1, h * w, attention_dim])
  if valid_mask is not None:
    valid_mask_q = tf.reshape(valid_mask, [-1, h * w, 1])

  # Adopt idea from Linformer (https://arxiv.org/abs/2006.04768) to reduce the
  # memory usage. Instead of learning a downsample function, we use determinstic
  # image downsample methods (nearest neighbor or bilinear) to reduce the size
  # of key and value.
  if linformer:
    k = resize_fn(
        k, ((h // factor + 1), (w // factor + 1)), align_corners=True)
    k = tf.reshape(k,
                   [-1, (h // factor + 1) * (w // factor + 1), attention_dim])

    if valid_mask is not None:
      valid_mask_k = tf.compat.v1.image.resize_nearest_neighbor(
          valid_mask, ((h // factor + 1), (w // factor + 1)))
      valid_mask_k = tf.reshape(
          tf.cast(valid_mask_k, tf.float32),
          [-1, (h // factor + 1) * (w // factor + 1), 1])
  else:
    k = tf.reshape(k, [-1, h * w, attention_dim])
    valid_mask_k = tf.reshape(valid_mask, [-1, h * w, 1])

  matmul_qk = tf.matmul(q, k, transpose_b=True)
  scaled_att_logits = matmul_qk / math.sqrt(attention_dim)
  # Masking
  if valid_mask is not None:
    final_mask = tf.matmul(valid_mask_q, valid_mask_k, transpose_b=True)
    scaled_att_logits += (1 - final_mask) * -1e9
  att_weights = tf.nn.softmax(scaled_att_logits, axis=-1)

  if linformer:
    value = resize_fn(
        score, ((h // factor + 1), (w // factor + 1)), align_corners=True)
    value = tf.reshape(value,
                       [-1, (h // factor + 1) * (w // factor + 1), value_dim])
  else:
    value = tf.reshape(score, [-1, h * w, value_dim])
  att_score = tf.matmul(att_weights, value)
  att_score = tf.reshape(att_score, tf.shape(score))

  ## Add skip-connection and 1x1 conv to convert score back to logit
  att_score += score
  if value_dim != num_class:
    # Set an initial score for the background class. Since the score range of a
    # class is [0, 2] after skip-connection, we use 2 minus the max class
    # probability to set the initial background score for each pixel.
    bg = 2 - tf.reduce_max(att_score, axis=3, keepdims=True)
    att_score = tf.concat([bg, att_score], axis=-1)

  out_att_logits = slim.conv2d(
      att_score,
      num_class, [1, 1],
      scope='pixel_normalization',
      activation_fn=None,
      normalizer_fn=slim.batch_norm,
      normalizer_params={'is_training': is_training},
      reuse=tf.AUTO_REUSE)

  return out_att_logits