Create Analogous Head in Keras

Question

When I create a CNN in Fast AI using transfer learning a head like this is created:

(1): Sequential(
    (0): AdaptiveConcatPool2d(
      (ap): AdaptiveAvgPool2d(output_size=1)
      (mp): AdaptiveMaxPool2d(output_size=1)
    )
    (1): Flatten()
    (2): BatchNorm1d(3840, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (3): Dropout(p=0.05005, inplace=False)
    (4): Linear(in_features=3840, out_features=512, bias=True)
    (5): ReLU(inplace=True)
    (6): BatchNorm1d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (7): Dropout(p=0.1001, inplace=False)
    (8): Linear(in_features=512, out_features=3, bias=True)

I would like to create one as close as possible to this one in Keras, however I am stuck on the AdaptiveConcatPool2D part of this. There doesn't appear to be any classes like this currently available. Any ideas on how to replicate this as closely as possible in Keras?

Answer 1

You can write your own AdaptiveConcatPool2d in tf.keras. The layers AdaptiveAveragePooling2D and AdaptiveMaxPooling2D (which make up AdaptiveConcatPool2d) are available in TensorFlow Addons, and the outputs of those would have to be concatenated. Also see the TensorFlow documentation on composing layers with other layers. Note that PyTorch (on which FastAI is based) uses the "channels_first" data format.

TensorFlow addons will have to be installed: python -m pip install --no-cache-dir tensorflow-addons

import tensorflow as tf
import tensorflow_addons as tfa

tfkl = tf.keras.layers

class AdaptiveConcatPool2D(tfkl.Layer):

    def __init__(self, output_size=None, data_format="channels_last"):
        super().__init__()
        self.output_size = output_size or 1
        self.data_format = data_format
        self.avg_layer = tfa.layers.AdaptiveAveragePooling2D(self.output_size, data_format=self.data_format)
        self.max_layer = tfa.layers.AdaptiveMaxPooling2D(self.output_size, data_format=self.data_format)

    def call(self, x):
        # See https://github.com/fastai/fastai/blob/23695330a9b8fe173b16013362047ca635e92ea4/fastai2/layers.py#L112-L118
        if self.data_format == "channels_last":
            return tfkl.Concatenate(axis=-1)([self.max_layer(x), self.avg_layer(x)])
        elif self.data_format == "channels_first":
            return tfkl.Concatenate(axis=1)([self.max_layer(x), self.avg_layer(x)])

---

Tests

import fastai.layers
import numpy as np
import torch

a = np.random.uniform(size=[2, 3, 16, 16]).astype(np.float32)

for output_size in [1, 2, 4, 8]:
    out_tf = AdaptiveConcatPool2D(output_size, data_format="channels_first")(a)
    out_fastai = fastai.layers.AdaptiveConcatPool2d(output_size)(torch.from_numpy(a))
    assert np.allclose(out_tf.numpy(), out_fastai.numpy())