# -*- coding: utf-8 -*-
# File: dropblock.py
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import re
import six
# from tensorpack.tfutils.compat import tfv1 as tf  # this should be avoided first in model code
from tensorpack.tfutils.tower import get_current_tower_context
from tensorpack.models import GlobalAvgPooling, FullyConnected
import tensorflow as tf

__all__ = ['dropblock', 'dropblock2','dropblock3','dropblock4'] # 1: paper baseline; 2: group dropout; 3: group soft-dropout; 4: Uout group dropout  

def dropblock(net, keep_prob, dropblock_size, gap_w=None, label=None, G=None, CG=None, data_format='channels_first'):
  """DropBlock: a regularization method for convolutional neural networks.
  DropBlock is a form of structured dropout, where units in a contiguous
  region of a feature map are dropped together. DropBlock works better than
  dropout on convolutional layers due to the fact that activation units in
  convolutional layers are spatially correlated.
  See https://arxiv.org/pdf/1810.12890.pdf for details.
  Args:
    net: `Tensor` input tensor.
    is_training: `bool` for whether the model is training.
    keep_prob: `float` or `Tensor` keep_prob parameter of DropBlock. "None"
        means no DropBlock.
    dropblock_size: `int` size of blocks to be dropped by DropBlock.
    data_format: `str` either "channels_first" for `[batch, channels, height,
        width]` or "channels_last for `[batch, height, width, channels]`.
  Returns:
      A version of input tensor with DropBlock applied.
  Raises:
      if width and height of the input tensor are not equal.
  """

  ctx = get_current_tower_context()
  is_training = bool(ctx.is_training)

  if not is_training or keep_prob is None:
    return net

  tf.logging.info('Applying DropBlock: dropblock_size {}, net.shape {}'.format(dropblock_size, net.shape))

  if data_format == 'channels_last':
    _, width, height, _ = net.get_shape().as_list()
  else:
    _, _, width, height = net.get_shape().as_list()
  if width != height:
    raise ValueError('Input tensor with width!=height is not supported.')

  dropblock_size = min(dropblock_size, width)
  # seed_drop_rate is the gamma parameter of DropBlcok.
  seed_drop_rate = (1.0 - keep_prob) * width**2 / dropblock_size**2 / (
      width - dropblock_size + 1)**2

  # Forces the block to be inside the feature map.
  w_i, h_i = tf.meshgrid(tf.range(width), tf.range(width))
  valid_block_center = tf.logical_and(
      tf.logical_and(w_i >= int(dropblock_size // 2),
                     w_i < width - (dropblock_size - 1) // 2),
      tf.logical_and(h_i >= int(dropblock_size // 2),
                     h_i < width - (dropblock_size - 1) // 2))

  valid_block_center = tf.expand_dims(valid_block_center, 0)
  valid_block_center = tf.expand_dims(
      valid_block_center, -1 if data_format == 'channels_last' else 0)

  randnoise = tf.random_uniform(tf.shape(net), dtype=tf.float32)
  block_pattern = (1 - tf.cast(valid_block_center, dtype=tf.float32) + tf.cast(
      (1 - seed_drop_rate), dtype=tf.float32) + randnoise) >= 1
  block_pattern = tf.cast(block_pattern, dtype=tf.float32)

  if dropblock_size == width:
    block_pattern = tf.reduce_min(
        block_pattern,
        axis=[1, 2] if data_format == 'channels_last' else [2, 3],
        keepdims=True)
  else:
    if data_format == 'channels_last':
      ksize = [1, dropblock_size, dropblock_size, 1]
    else:
      ksize = [1, 1, dropblock_size, dropblock_size]
    block_pattern = -tf.nn.max_pool(
        -block_pattern, ksize=ksize, strides=[1, 1, 1, 1], padding='SAME',
        data_format='NHWC' if data_format == 'channels_last' else 'NCHW')

  percent_ones = tf.cast(tf.reduce_sum((block_pattern)), tf.float32) / tf.cast(
      tf.size(block_pattern), tf.float32)

  net = net / tf.cast(percent_ones, net.dtype) * tf.cast(
      block_pattern, net.dtype)
  return net


def dropblock2(net, keep_prob, dropblock_size, G=None, CG=None, data_format='channels_first'):
  """
  mimic GN
  """  
  ctx = get_current_tower_context()
  is_training = bool(ctx.is_training)

  if not is_training or keep_prob is None:
    return net

  tf.logging.info('Applying DropBlock: dropblock_size {}, net.shape {}'.format(dropblock_size, net.shape))

  if data_format == 'channels_last':
    N, height, width, C = net.get_shape().as_list()
  else:
    N, C, height, width = net.get_shape().as_list()
  N = tf.shape(net)[0] 
  if width != height:
    raise ValueError('Input tensor with width!=height is not supported.')

  if G == None: G = C // CG
  if CG == None: CG = C // G
  net = tf.reshape(net, [N, G, CG, height, width]) 

  dropblock_size = min(dropblock_size, width)
  # seed_drop_rate is the gamma parameter of DropBlcok.
  # seed_drop_rate = (1.0 - keep_prob) * width**2 * G**2 / (C * dropblock_size**2) / (C * (width - dropblock_size + 1)**2)
  seed_drop_rate = (1.0 - keep_prob) * width**2 / dropblock_size**2 / (width - dropblock_size + 1)**2

  # Forces the block to be inside the feature map.
  w_i, h_i = tf.meshgrid(tf.range(width), tf.range(width))
  valid_block_center = tf.logical_and(
      tf.logical_and(w_i >= int(dropblock_size // 2),
                     w_i < width - (dropblock_size - 1) // 2),
      tf.logical_and(h_i >= int(dropblock_size // 2),
                     h_i < width - (dropblock_size - 1) // 2))

  valid_block_center = tf.expand_dims(valid_block_center, 0) # for depth
  valid_block_center = tf.expand_dims(valid_block_center, 0) # for batch
  valid_block_center = tf.expand_dims(valid_block_center, 0) # for channel
  randnoise = tf.random_uniform([N, G, 1, width, height], dtype=tf.float32)

  block_pattern = (1 - tf.cast(valid_block_center, dtype=tf.float32) + tf.cast(
      (1 - seed_drop_rate), dtype=tf.float32) + randnoise) >= 1
  block_pattern = tf.cast(block_pattern, dtype=tf.float32)

  if dropblock_size == width:
    block_pattern = tf.reduce_min(block_pattern, axis=[2, 3, 4], keepdims=True)
  else:
    ksize = [1, 1, dropblock_size, dropblock_size]
    
    block_pattern = tf.reduce_max(-block_pattern, reduction_indices=[2])
    block_pattern = -tf.nn.max_pool(block_pattern, ksize=ksize, strides=[1, 1, 1, 1], padding='SAME', data_format='NCHW')
    block_pattern = tf.expand_dims(block_pattern, 2)

  percent_ones = tf.cast(tf.reduce_sum((block_pattern)), tf.float32) / tf.cast(tf.size(block_pattern), tf.float32)
  net = net / tf.cast(percent_ones, net.dtype) * tf.cast(block_pattern, net.dtype)
  net = tf.reshape(net, [N, height, width, C]) if data_format == 'channels_last' else tf.reshape(net, [N, C, height, width])
  return net


def CamDrop(net, keep_prob, dropblock_size, flag=None, label=None, G=None, CG=None, data_format='channels_first'):
  '''CamDrop'''
  def _get_cam(net, label, flag, dropblock_size, data_format='channels_first'):
    '''
    net: [N, C, H, W]
    gap_w : [gap_C, num_of_class]
    '''
    if data_format == 'channels_last':
      N, height, width, C = net.get_shape().as_list()
    else:
      N, C, height, width = net.get_shape().as_list()
    N = tf.shape(net)[0] 

    gap_w = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, 'linear/W') if flag > 0 else None

    if not gap_w is None:
      gap_w = tf.convert_to_tensor(gap_w, tf.float32)
      gap_C, num = tf.squeeze(gap_w, 0).get_shape().as_list() # [gap_C, num]
      gap_w = tf.reshape(gap_w, [C, gap_C//C, num])
      gap_w = tf.reduce_mean(gap_w, reduction_indices=[1]) # [C, num]
      label = tf.gather(tf.transpose(gap_w), label) # [N, C]

      # spatial 
      weights = tf.expand_dims(label, 2) # [N, C, 1]
      net = tf.reshape(net, [N, height*width, C]) if data_format == 'channels_last' else tf.reshape(net, [N, C, height*width])
      cam = tf.matmul(weights, net, transpose_a=True) # [N, 1, width*height]
      # spt_mask = tf.not_equal(cam, tf.reduce_max(cam, reduction_indices=[2], keepdims=True))

      # cam = tf.reshape(cam, [N, height, width, 1]) if data_format == 'channels_last' else tf.reshape(cam, [N, 1, height, width])
      # cam = tf.nn.avg_pool(cam, ksize=[1, 1, dropblock_size, dropblock_size], strides=[1, 1, 1, 1], padding='VALID', data_format='NCHW')
      # left_or_top = (dropblock_size-1) // 2 
      # right_or_bot = left_or_top if dropblock_size % 2 == 1 else dropblock_size-left_or_top-1
      # cam = tf.pad(cam, [[0, 0], [0, 0], [left_or_top, right_or_bot], [left_or_top, right_or_bot]])
      # cam = tf.reshape(cam, [N, height*width, 1]) if data_format == 'channels_last' else tf.reshape(cam, [N, 1, height*width])


      k = tf.cast(height*width/dropblock_size**2, tf.int32)
      topk, _ = tf.math.top_k(cam, k=k) # [N, 1, k]
      topk = tf.gather(topk, indices=[k-1], axis=-1) # [N, 1, 1]
      spt_mask = (cam < topk)

      spt_mask = tf.reshape(spt_mask, [N, height, width, 1]) if data_format == 'channels_last' else tf.reshape(spt_mask, [N, 1, height, width])

      # channel
      k = tf.cast(C/8, tf.int32)
      topk, _ = tf.math.top_k(label, k=k+1) # [N, k]
      topk = tf.gather(topk, indices=k, axis=1) # [N, 1]
      topk = tf.expand_dims(topk, 1) # [N, C, 1]
      chan_mask = (label < topk)
      chan_mask = tf.expand_dims(chan_mask, 2) # [N, C, 1]
      chan_mask = tf.expand_dims(chan_mask, 2) # [N, C, 1, 1]

      cam_mask = tf.logical_or(spt_mask, chan_mask)
      # chan_mask = tf.reshape(tf.nn.softmax(cam), [N*C, height*width]) if data_format == 'channels_last' else tf.reshape(tf.nn.softmax(cam), [N*C, height*width])
      # chan_mask = tf.reshape(cam, [N*C, height*width]) if data_format == 'channels_last' else tf.reshape(cam, [N*C, height*width])
      # chan_mask = tf.reshape(tf.nn.sigmoid(cam), [N, height, width, 1]) if data_format == 'channels_last' else tf.reshape(tf.nn.sigmoid(cam), [N, 1, height, width])
    else:
      cam_mask = False
    return cam_mask

  # def _get_gradcam(net, cost=None, gap_w=None, data_format='channels_first'):
  #   # Conv layer tensor [?,2048,10,10]
  #   def _compute_gradients(tensor, var_list):
  #     grads = tf.gradients(tensor, var_list)
  #     return [grad if grad is not None else tf.zeros_like(var)
  #             for var, grad in zip(var_list, grads)]
  #   # grads = tf.gradients(cost, net)[0]
  #   if not gap_w is None:
  #     # Normalizing the gradients
  #     if data_format == 'channels_last':
  #       N, height, width, C = net.get_shape().as_list()
  #     else:
  #       N, C, height, width = net.get_shape().as_list()
  #     N = tf.shape(net)[0] 

  #     grads = _compute_gradients(cost, [net])[0]
  #     norm_grads = tf.divide(grads, tf.sqrt(tf.reduce_mean(tf.square(grads), reduction_indices=[2,3], keepdims=True)) + tf.constant(1e-5))
  #     weights = tf.reduce_mean(norm_grads, reduction_indices=[2,3])  # [N, C]
  #     weights = tf.expand_dims(weights, 2) # [N, C, 1]
  #     net = tf.reshape(net, [N, height*width, C]) if data_format == 'channels_last' else tf.reshape(net, [N, C, height*width])
  #     # cam_mean = 1 + tf.matmul(net, weights, transpose_a=True) # [N, width*height, 1]
  #     cam_mean = tf.maximum(tf.matmul(weights, net, transpose_a=True), 0) # [N, 1, width*height]
  #     cam_chan = tf.maximum(tf.multiply(net, weights), 0) # [N, C, width*height]
  #     cam = cam_mean*cam_chan
  #     # Passing through ReLU
  #     cam = cam / tf.reduce_max(cam, reduction_indices=[1,2], keepdims=True)
  #     cam = tf.reshape(cam, [N, height, width, C]) if data_format == 'channels_last' else tf.reshape(cam, [N, C, height, width])

  #   else:
  #     cam = 0.
  #   return cam

  # def _gumbel_softmax(logits, tau, shape, seed_drop_rate, eps=1e-20):
  #   if logits == False:
  #     return logits

  #   U = tf.random_uniform(tf.shape(logits), minval=0, maxval=1)
  #   y = logits - tf.log(-tf.log(U + eps) + eps)
  #   cam_mask = tf.nn.softmax(y / tau)
  #   topk, _ = tf.math.top_k(cam_mask, k=tf.cast(seed_drop_rate*shape[-1], tf.int32)) # [N, 1]
  #   topk = tf.gather(topk, indices=tf.cast(seed_drop_rate*shape[-1], tf.int32)-1, axis=1)
  #   topk = tf.expand_dims(topk, 1) # [N, C, 1]
  #   cam_mask = (cam_mask < topk)
  #   # cam_mask = tf.cast(tf.equal(cam_mask, tf.reduce_max(cam_mask, reduction_indices=[1], keepdims=True)), tf.float32)

  #   cam_mask = tf.expand_dims(cam_mask, 2) # [N, C, 1]
  #   cam_mask = tf.expand_dims(cam_mask, 2) # [N, C, 1, 1]
  #   return cam_mask


  ctx = get_current_tower_context()
  is_training = bool(ctx.is_training)

  if not is_training or keep_prob is None:
    return net

  tf.logging.info('Applying DropBlock: dropblock_size {}, net.shape {}'.format(dropblock_size, net.shape))

  if data_format == 'channels_last':
    _, width, height, C = net.get_shape().as_list()
  else:
    _, C, width, height = net.get_shape().as_list()
  if width != height:
    raise ValueError('Input tensor with width!=height is not supported.')
  N = tf.shape(net)[0] 

  dropblock_size = min(dropblock_size, width)
  # seed_drop_rate is the gamma parameter of DropBlcok.
  seed_drop_rate = (1.0 - keep_prob) * width**2 / dropblock_size**2 / (width - dropblock_size + 1)**2

  cam_mask = _get_cam(net, label, flag, dropblock_size, data_format)

  # Forces the block to be inside the feature map.
  w_i, h_i = tf.meshgrid(tf.range(width), tf.range(width))
  valid_block_center = tf.logical_and(
      tf.logical_and(w_i >= int(dropblock_size // 2),
                     w_i < width - (dropblock_size - 1) // 2),
      tf.logical_and(h_i >= int(dropblock_size // 2),
                     h_i < width - (dropblock_size - 1) // 2))

  valid_block_center = tf.expand_dims(valid_block_center, 0)
  valid_block_center = tf.expand_dims(valid_block_center, -1 if data_format == 'channels_last' else 0)

  randnoise = tf.random_uniform(tf.shape(net), dtype=tf.float32)
  block_pattern = (1 - tf.cast(valid_block_center, dtype=tf.float32) + tf.cast((1 - seed_drop_rate), dtype=tf.float32) + randnoise) >= 1
  block_pattern = tf.logical_or(block_pattern, cam_mask)
  block_pattern = tf.cast(block_pattern, dtype=tf.float32)

  if dropblock_size == width:
    block_pattern = tf.reduce_min(
        block_pattern,
        axis=[1, 2] if data_format == 'channels_last' else [2, 3],
        keepdims=True)
  else:
    if data_format == 'channels_last':
      ksize = [1, dropblock_size, dropblock_size, 1]
    else:
      ksize = [1, 1, dropblock_size, dropblock_size]
    block_pattern = -tf.nn.max_pool(
        -block_pattern, ksize=ksize, strides=[1, 1, 1, 1], padding='SAME',
        data_format='NHWC' if data_format == 'channels_last' else 'NCHW')

  percent_ones = tf.cast(tf.reduce_sum((block_pattern)), tf.float32) / tf.cast(tf.size(block_pattern), tf.float32)

  net = net / tf.cast(percent_ones, net.dtype) * tf.cast(block_pattern, net.dtype)
  return net
[((37, 8, 37, 35), 'tensorpack.tfutils.tower.get_current_tower_context', 'get_current_tower_context', ({}, {}), '()', False, 'from tensorpack.tfutils.tower import get_current_tower_context\n'), ((65, 23, 65, 60), 'tensorflow.expand_dims', 'tf.expand_dims', ({(65, 38, 65, 56): 'valid_block_center', (65, 58, 65, 59): '0'}, {}), '(valid_block_center, 0)', True, 'import tensorflow as tf\n'), ((66, 23, 67, 70), 'tensorflow.expand_dims', 'tf.expand_dims', ({(67, 6, 67, 24): 'valid_block_center', (67, 26, 67, 69): "-1 if data_format == 'channels_last' else 0"}, {}), "(valid_block_center, -1 if data_format == 'channels_last' else 0)", True, 'import tensorflow as tf\n'), ((72, 18, 72, 58), 'tensorflow.cast', 'tf.cast', (), '', True, 'import tensorflow as tf\n'), ((100, 8, 100, 35), 'tensorpack.tfutils.tower.get_current_tower_context', 'get_current_tower_context', ({}, {}), '()', False, 'from tensorpack.tfutils.tower import get_current_tower_context\n'), ((118, 8, 118, 50), 'tensorflow.reshape', 'tf.reshape', ({(118, 19, 118, 22): 'net', (118, 24, 118, 49): '[N, G, CG, height, width]'}, {}), '(net, [N, G, CG, height, width])', True, 'import tensorflow as tf\n'), ((133, 23, 133, 60), 'tensorflow.expand_dims', 'tf.expand_dims', ({(133, 38, 133, 56): 'valid_block_center', (133, 58, 133, 59): '0'}, {}), '(valid_block_center, 0)', True, 'import tensorflow as tf\n'), ((134, 23, 134, 60), 'tensorflow.expand_dims', 'tf.expand_dims', ({(134, 38, 134, 56): 'valid_block_center', (134, 58, 134, 59): '0'}, {}), '(valid_block_center, 0)', True, 'import tensorflow as tf\n'), ((135, 23, 135, 60), 'tensorflow.expand_dims', 'tf.expand_dims', ({(135, 38, 135, 56): 'valid_block_center', (135, 58, 135, 59): '0'}, {}), '(valid_block_center, 0)', True, 'import tensorflow as tf\n'), ((136, 14, 136, 75), 'tensorflow.random_uniform', 'tf.random_uniform', (), '', True, 'import tensorflow as tf\n'), ((140, 18, 140, 58), 'tensorflow.cast', 'tf.cast', (), '', True, 'import tensorflow as tf\n'), ((267, 8, 267, 35), 'tensorpack.tfutils.tower.get_current_tower_context', 'get_current_tower_context', ({}, {}), '()', False, 'from tensorpack.tfutils.tower import get_current_tower_context\n'), ((297, 23, 297, 60), 'tensorflow.expand_dims', 'tf.expand_dims', ({(297, 38, 297, 56): 'valid_block_center', (297, 58, 297, 59): '0'}, {}), '(valid_block_center, 0)', True, 'import tensorflow as tf\n'), ((298, 23, 298, 102), 'tensorflow.expand_dims', 'tf.expand_dims', ({(298, 38, 298, 56): 'valid_block_center', (298, 58, 298, 101): "-1 if data_format == 'channels_last' else 0"}, {}), "(valid_block_center, -1 if data_format == 'channels_last' else 0)", True, 'import tensorflow as tf\n'), ((302, 18, 302, 56), 'tensorflow.logical_or', 'tf.logical_or', ({(302, 32, 302, 45): 'block_pattern', (302, 47, 302, 55): 'cam_mask'}, {}), '(block_pattern, cam_mask)', True, 'import tensorflow as tf\n'), ((303, 18, 303, 58), 'tensorflow.cast', 'tf.cast', (), '', True, 'import tensorflow as tf\n'), ((58, 25, 58, 40), 'tensorflow.range', 'tf.range', ({(58, 34, 58, 39): 'width'}, {}), '(width)', True, 'import tensorflow as tf\n'), ((58, 42, 58, 57), 'tensorflow.range', 'tf.range', ({(58, 51, 58, 56): 'width'}, {}), '(width)', True, 'import tensorflow as tf\n'), ((69, 32, 69, 45), 'tensorflow.shape', 'tf.shape', ({(69, 41, 69, 44): 'net'}, {}), '(net)', True, 'import tensorflow as tf\n'), ((75, 20, 78, 22), 'tensorflow.reduce_min', 'tf.reduce_min', (), '', True, 'import tensorflow as tf\n'), ((91, 49, 92, 31), 'tensorflow.cast', 'tf.cast', ({(92, 6, 92, 19): 'block_pattern', (92, 21, 92, 30): 'net.dtype'}, {}), '(block_pattern, net.dtype)', True, 'import tensorflow as tf\n'), ((112, 6, 112, 19), 'tensorflow.shape', 'tf.shape', ({(112, 15, 112, 18): 'net'}, {}), '(net)', True, 'import tensorflow as tf\n'), ((126, 25, 126, 40), 'tensorflow.range', 'tf.range', ({(126, 34, 126, 39): 'width'}, {}), '(width)', True, 'import tensorflow as tf\n'), ((126, 42, 126, 57), 'tensorflow.range', 'tf.range', ({(126, 51, 126, 56): 'width'}, {}), '(width)', True, 'import tensorflow as tf\n'), ((143, 20, 143, 79), 'tensorflow.reduce_min', 'tf.reduce_min', (), '', True, 'import tensorflow as tf\n'), ((147, 20, 147, 72), 'tensorflow.reduce_max', 'tf.reduce_max', (), '', True, 'import tensorflow as tf\n'), ((149, 20, 149, 52), 'tensorflow.expand_dims', 'tf.expand_dims', ({(149, 35, 149, 48): 'block_pattern', (149, 50, 149, 51): '2'}, {}), '(block_pattern, 2)', True, 'import tensorflow as tf\n'), ((152, 49, 152, 82), 'tensorflow.cast', 'tf.cast', ({(152, 57, 152, 70): 'block_pattern', (152, 72, 152, 81): 'net.dtype'}, {}), '(block_pattern, net.dtype)', True, 'import tensorflow as tf\n'), ((153, 8, 153, 46), 'tensorflow.reshape', 'tf.reshape', ({(153, 19, 153, 22): 'net', (153, 24, 153, 45): '[N, height, width, C]'}, {}), '(net, [N, height, width, C])', True, 'import tensorflow as tf\n'), ((153, 86, 153, 124), 'tensorflow.reshape', 'tf.reshape', ({(153, 97, 153, 100): 'net', (153, 102, 153, 123): '[N, C, height, width]'}, {}), '(net, [N, C, height, width])', True, 'import tensorflow as tf\n'), ((281, 6, 281, 19), 'tensorflow.shape', 'tf.shape', ({(281, 15, 281, 18): 'net'}, {}), '(net)', True, 'import tensorflow as tf\n'), ((290, 25, 290, 40), 'tensorflow.range', 'tf.range', ({(290, 34, 290, 39): 'width'}, {}), '(width)', True, 'import tensorflow as tf\n'), ((290, 42, 290, 57), 'tensorflow.range', 'tf.range', ({(290, 51, 290, 56): 'width'}, {}), '(width)', True, 'import tensorflow as tf\n'), ((300, 32, 300, 45), 'tensorflow.shape', 'tf.shape', ({(300, 41, 300, 44): 'net'}, {}), '(net)', True, 'import tensorflow as tf\n'), ((306, 20, 309, 22), 'tensorflow.reduce_min', 'tf.reduce_min', (), '', True, 'import tensorflow as tf\n'), ((321, 49, 321, 82), 'tensorflow.cast', 'tf.cast', ({(321, 57, 321, 70): 'block_pattern', (321, 72, 321, 81): 'net.dtype'}, {}), '(block_pattern, net.dtype)', True, 'import tensorflow as tf\n'), ((84, 21, 86, 73), 'tensorflow.nn.max_pool', 'tf.nn.max_pool', (), '', True, 'import tensorflow as tf\n'), ((88, 25, 88, 55), 'tensorflow.reduce_sum', 'tf.reduce_sum', ({(88, 40, 88, 53): 'block_pattern'}, {}), '(block_pattern)', True, 'import tensorflow as tf\n'), ((89, 6, 89, 28), 'tensorflow.size', 'tf.size', ({(89, 14, 89, 27): 'block_pattern'}, {}), '(block_pattern)', True, 'import tensorflow as tf\n'), ((91, 14, 91, 46), 'tensorflow.cast', 'tf.cast', ({(91, 22, 91, 34): 'percent_ones', (91, 36, 91, 45): 'net.dtype'}, {}), '(percent_ones, net.dtype)', True, 'import tensorflow as tf\n'), ((148, 21, 148, 121), 'tensorflow.nn.max_pool', 'tf.nn.max_pool', (), '', True, 'import tensorflow as tf\n'), ((151, 25, 151, 55), 'tensorflow.reduce_sum', 'tf.reduce_sum', ({(151, 40, 151, 53): 'block_pattern'}, {}), '(block_pattern)', True, 'import tensorflow as tf\n'), ((151, 79, 151, 101), 'tensorflow.size', 'tf.size', ({(151, 87, 151, 100): 'block_pattern'}, {}), '(block_pattern)', True, 'import tensorflow as tf\n'), ((152, 14, 152, 46), 'tensorflow.cast', 'tf.cast', ({(152, 22, 152, 34): 'percent_ones', (152, 36, 152, 45): 'net.dtype'}, {}), '(percent_ones, net.dtype)', True, 'import tensorflow as tf\n'), ((168, 8, 168, 21), 'tensorflow.shape', 'tf.shape', ({(168, 17, 168, 20): 'net'}, {}), '(net)', True, 'import tensorflow as tf\n'), ((170, 12, 170, 75), 'tensorflow.get_collection', 'tf.get_collection', ({(170, 30, 170, 62): 'tf.GraphKeys.TRAINABLE_VARIABLES', (170, 64, 170, 74): '"""linear/W"""'}, {}), "(tf.GraphKeys.TRAINABLE_VARIABLES, 'linear/W')", True, 'import tensorflow as tf\n'), ((173, 14, 173, 53), 'tensorflow.convert_to_tensor', 'tf.convert_to_tensor', ({(173, 35, 173, 40): 'gap_w', (173, 42, 173, 52): 'tf.float32'}, {}), '(gap_w, tf.float32)', True, 'import tensorflow as tf\n'), ((175, 14, 175, 51), 'tensorflow.reshape', 'tf.reshape', ({(175, 25, 175, 30): 'gap_w', (175, 32, 175, 50): '[C, gap_C // C, num]'}, {}), '(gap_w, [C, gap_C // C, num])', True, 'import tensorflow as tf\n'), ((176, 14, 176, 58), 'tensorflow.reduce_mean', 'tf.reduce_mean', (), '', True, 'import tensorflow as tf\n'), ((180, 16, 180, 40), 'tensorflow.expand_dims', 'tf.expand_dims', ({(180, 31, 180, 36): 'label', (180, 38, 180, 39): '2'}, {}), '(label, 2)', True, 'import tensorflow as tf\n'), ((182, 12, 182, 53), 'tensorflow.matmul', 'tf.matmul', (), '', True, 'import tensorflow as tf\n'), ((193, 10, 193, 59), 'tensorflow.cast', 'tf.cast', ({(193, 18, 193, 48): 'height * width / dropblock_size ** 2', (193, 50, 193, 58): 'tf.int32'}, {}), '(height * width / dropblock_size ** 2, tf.int32)', True, 'import tensorflow as tf\n'), ((194, 16, 194, 39), 'tensorflow.math.top_k', 'tf.math.top_k', (), '', True, 'import tensorflow as tf\n'), ((195, 13, 195, 52), 'tensorflow.gather', 'tf.gather', (), '', True, 'import tensorflow as tf\n'), ((201, 10, 201, 32), 'tensorflow.cast', 'tf.cast', ({(201, 18, 201, 21): 'C / 8', (201, 23, 201, 31): 'tf.int32'}, {}), '(C / 8, tf.int32)', True, 'import tensorflow as tf\n'), ((202, 16, 202, 43), 'tensorflow.math.top_k', 'tf.math.top_k', (), '', True, 'import tensorflow as tf\n'), ((203, 13, 203, 47), 'tensorflow.gather', 'tf.gather', (), '', True, 'import tensorflow as tf\n'), ((204, 13, 204, 36), 'tensorflow.expand_dims', 'tf.expand_dims', ({(204, 28, 204, 32): 'topk', (204, 34, 204, 35): '1'}, {}), '(topk, 1)', True, 'import tensorflow as tf\n'), ((206, 18, 206, 46), 'tensorflow.expand_dims', 'tf.expand_dims', ({(206, 33, 206, 42): 'chan_mask', (206, 44, 206, 45): '2'}, {}), '(chan_mask, 2)', True, 'import tensorflow as tf\n'), ((207, 18, 207, 46), 'tensorflow.expand_dims', 'tf.expand_dims', ({(207, 33, 207, 42): 'chan_mask', (207, 44, 207, 45): '2'}, {}), '(chan_mask, 2)', True, 'import tensorflow as tf\n'), ((209, 17, 209, 51), 'tensorflow.logical_or', 'tf.logical_or', ({(209, 31, 209, 39): 'spt_mask', (209, 41, 209, 50): 'chan_mask'}, {}), '(spt_mask, chan_mask)', True, 'import tensorflow as tf\n'), ((315, 21, 317, 73), 'tensorflow.nn.max_pool', 'tf.nn.max_pool', (), '', True, 'import tensorflow as tf\n'), ((319, 25, 319, 55), 'tensorflow.reduce_sum', 'tf.reduce_sum', ({(319, 40, 319, 53): 'block_pattern'}, {}), '(block_pattern)', True, 'import tensorflow as tf\n'), ((319, 79, 319, 101), 'tensorflow.size', 'tf.size', ({(319, 87, 319, 100): 'block_pattern'}, {}), '(block_pattern)', True, 'import tensorflow as tf\n'), ((321, 14, 321, 46), 'tensorflow.cast', 'tf.cast', ({(321, 22, 321, 34): 'percent_ones', (321, 36, 321, 45): 'net.dtype'}, {}), '(percent_ones, net.dtype)', True, 'import tensorflow as tf\n'), ((70, 71, 71, 45), 'tensorflow.cast', 'tf.cast', (), '', True, 'import tensorflow as tf\n'), ((138, 71, 139, 45), 'tensorflow.cast', 'tf.cast', (), '', True, 'import tensorflow as tf\n'), ((177, 24, 177, 43), 'tensorflow.transpose', 'tf.transpose', ({(177, 37, 177, 42): 'gap_w'}, {}), '(gap_w)', True, 'import tensorflow as tf\n'), ((181, 12, 181, 49), 'tensorflow.reshape', 'tf.reshape', ({(181, 23, 181, 26): 'net', (181, 28, 181, 48): '[N, height * width, C]'}, {}), '(net, [N, height * width, C])', True, 'import tensorflow as tf\n'), ((181, 89, 181, 126), 'tensorflow.reshape', 'tf.reshape', ({(181, 100, 181, 103): 'net', (181, 105, 181, 125): '[N, C, height * width]'}, {}), '(net, [N, C, height * width])', True, 'import tensorflow as tf\n'), ((198, 17, 198, 60), 'tensorflow.reshape', 'tf.reshape', ({(198, 28, 198, 36): 'spt_mask', (198, 38, 198, 59): '[N, height, width, 1]'}, {}), '(spt_mask, [N, height, width, 1])', True, 'import tensorflow as tf\n'), ((198, 100, 198, 143), 'tensorflow.reshape', 'tf.reshape', ({(198, 111, 198, 119): 'spt_mask', (198, 121, 198, 142): '[N, 1, height, width]'}, {}), '(spt_mask, [N, 1, height, width])', True, 'import tensorflow as tf\n'), ((301, 71, 301, 118), 'tensorflow.cast', 'tf.cast', (), '', True, 'import tensorflow as tf\n'), ((70, 23, 70, 68), 'tensorflow.cast', 'tf.cast', (), '', True, 'import tensorflow as tf\n'), ((138, 23, 138, 68), 'tensorflow.cast', 'tf.cast', (), '', True, 'import tensorflow as tf\n'), ((301, 23, 301, 68), 'tensorflow.cast', 'tf.cast', (), '', True, 'import tensorflow as tf\n'), ((174, 19, 174, 39), 'tensorflow.squeeze', 'tf.squeeze', ({(174, 30, 174, 35): 'gap_w', (174, 37, 174, 38): '0'}, {}), '(gap_w, 0)', True, 'import tensorflow as tf\n')]