How to get Attention Maps from Attention Gates in Attention UNET?

Question

Contex

I have Attention UNET for image segmentation. I use it for humans segmentation.

Question

Everything works fine. I want to get attention maps from my network, so I could see what my UNET is focused on.

Code

def conv_block(input, filter_num=32, kernel=(3, 3), max_pool=True, dropout=0.1):
  c1 = tf.keras.layers.Conv2D(filter_num, kernel, activation='relu', kernel_initializer='he_normal', padding='same')(input)
  if dropout:
    c1 = tf.keras.layers.Dropout(dropout)(c1)
  c1 = tf.keras.layers.Conv2D(filter_num, kernel, activation='relu', kernel_initializer='he_normal', padding='same')(c1)
  if max_pool:
    p1 = tf.keras.layers.MaxPooling2D((2, 2))(c1)
    return c1, p1
  else:
    return c1

def expend_as(tensor, rep):

    # Anonymous lambda function to expand the specified axis by a factor of argument, rep.
    # If tensor has shape (512,512,N), lambda will return a tensor of shape (512,512,N*rep), if specified axis=2

    my_repeat = tf.keras.layers.Lambda(lambda x, repnum: K.repeat_elements(x, repnum, axis=3),
                       arguments={'repnum': rep})(tensor)
    return my_repeat


def AttnGatingBlock(x, g, inter_shape):

    shape_x = K.int_shape(x)
    shape_g = K.int_shape(g)

    # Getting the gating signal to the same number of filters as the inter_shape
    phi_g = tf.keras.layers.Conv2D(filters=inter_shape,
                   kernel_size=1,
                   strides=1,
                   padding='same')(g)

    # Getting the x signal to the same shape as the gating signal
    theta_x = tf.keras.layers.Conv2D(filters=inter_shape,
                     kernel_size=3,
                     strides=(shape_x[1] // shape_g[1],
                              shape_x[2] // shape_g[2]),
                     padding='same')(x)

    # Element-wise addition of the gating and x signals
    add_xg = tf.keras.layers.add([phi_g, theta_x])
    add_xg = tf.keras.layers.Activation('relu')(add_xg)

    # 1x1x1 convolution
    psi = tf.keras.layers.Conv2D(filters=1, kernel_size=1, padding='same')(add_xg)
    psi = tf.keras.layers.Activation('sigmoid')(psi)
    shape_sigmoid = K.int_shape(psi)

    # Upsampling psi back to the original dimensions of x signal
    upsample_sigmoid_xg = tf.keras.layers.UpSampling2D(size=(shape_x[1] // shape_sigmoid[1],
                                             shape_x[2] //
                                             shape_sigmoid[2]))(psi)

    # Expanding the filter axis to the number of filters in the original x signal
    upsample_sigmoid_xg = expend_as(upsample_sigmoid_xg, shape_x[3])

    # Element-wise multiplication of attention coefficients back onto original x signal
    attn_coefficients = tf.keras.layers.multiply([upsample_sigmoid_xg, x])

    # Final 1x1x1 convolution to consolidate attention signal to original x dimensions
    output = tf.keras.layers.Conv2D(filters=shape_x[3],
                    kernel_size=1,
                    strides=1,
                    padding='same')(attn_coefficients)
    output = tf.keras.layers.BatchNormalization()(output)
    return output

with tpu_strategy.scope():
  inputs = tf.keras.layers.Input((IMG_HEIGHT, IMG_WIDTH, IMG_CHANNELS))
  s = tf.keras.layers.Lambda(lambda x: x / 255)(inputs)

  c1, p1 = conv_block(s, 32, (3, 3), True, 0.1)
  c2, p2 = conv_block(p1, 64, (3, 3), True, 0.1)
  c3, p3 = conv_block(p2, 128, (3, 3), True, 0.1)
  c4, p4 = conv_block(p3, 256, (3, 3), True, 0.1)
  c5 = conv_block(p4, 512, (3, 3), False, 0.1)

  u6 = tf.keras.layers.Conv2DTranspose(128, (2, 2), strides=(2, 2), padding='same')(c5)
  att1 = AttnGatingBlock(c4, c5, 256)
  u6 = tf.keras.layers.concatenate([u6, att1])
  c6 = conv_block(u6, 256, (3, 3), False, 0.1)

  u7 = tf.keras.layers.Conv2DTranspose(64, (2, 2), strides=(2, 2), padding='same')(c6)
  att2 = AttnGatingBlock(c3, c6, 128)
  u7 = tf.keras.layers.concatenate([u7, att2])
  c7 = conv_block(u7, 128, (3, 3), False, 0.1)

  u8 = tf.keras.layers.Conv2DTranspose(32, (2, 2), strides=(2, 2), padding='same')(c7)
  att3 = AttnGatingBlock(c2, c7, 64)
  u8 = tf.keras.layers.concatenate([u8, att3])
  c8 = conv_block(u8, 64, (3, 3), False, 0.1)

  u9 = tf.keras.layers.Conv2DTranspose(16, (2, 2), strides=(2, 2), padding='same')(c8)
  att4 = AttnGatingBlock(c1, c8, 32)
  u9 = tf.keras.layers.concatenate([u9, att4], axis=3)
  c9 = conv_block(u9, 32, (3, 3), False, 0.1)

  outputs = tf.keras.layers.Conv2D(1, (1, 1), activation='sigmoid')(c9)

  model = tf.keras.Model(inputs=[inputs], outputs=[outputs])
  model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
  plot_model(model)
  model.summary()
  callback = tf.keras.callbacks.EarlyStopping(monitor='val_loss', min_delta=0.01, patience=10, restore_best_weights=True)
  history = model.fit(train_X, train_y, batch_size=128, epochs=150, validation_split=0.1)

Link to whole code

https://colab.research.google.com/drive/1JELaVaVoeHu6Xfdh2P5hVTQR5ORn0uvS?usp=sharing

I'm editing my previous notebook with standard UNET so text cells might be a little bit not accurate. However dataset that I'm using is the same as in the link.