CEncoderHead.py

# a bit hacky way to make encoder head obtain input shape dynamically
import tensorflow as tf
import tensorflow.keras.layers as L
from NN.utils import sMLP
from NN.encoding import CCoordsGridLayer, CCoordsEncodingLayer
from NN.layers import MixerConvLayer, Patches, TransformerBlock
from Utils.utils import dumb_deepcopy

def block_params_from_config(config):
  layers = config.get('layers', None)
  if not(layers is None): return layers

  defaultConvParams = {
    'kernel size': config.get('kernel size', 3),
    'activation': config.get('activation', 'relu'),
    'name': config.get('name', 'Conv2D'),
  }
  convBefore = config['conv before']
  # if convBefore is an integer, then it's the same for all layers
  if isinstance(convBefore, int):
    convParams = { 'channels': config['channels'], **defaultConvParams }
    convBefore = [convParams] * convBefore # repeat the same parameters
    pass
  assert isinstance(convBefore, list), 'convBefore must be a list'
  # if convBefore is a list of integers, then each integer is the number of channels
  if (0 < len(convBefore)) and isinstance(convBefore[0], int):
    convBefore = [ {'channels': sz, **defaultConvParams} for sz in convBefore ]
    pass

  # add separately last layer
  lastConvParams = {
    'channels': config.get('channels last', config['channels']),
    'kernel size': config.get('kernel size last', defaultConvParams['kernel size']),
    'activation': config.get('final activation', defaultConvParams['activation']),
    'name': config.get('last name', 'Conv2D'),
  }
  return convBefore + [lastConvParams]
  
def conv_block_from_config(data, config, defaults, name='CB'):
  config = {**defaults, **config} # merge defaults and config
  convParams = block_params_from_config(config)
  # apply convolutions to the data
  for i, parameters in enumerate(convParams):
    parameters = dumb_deepcopy(parameters)
    Name = parameters.get('name', 'Conv2D')
    if 'Conv2D' == Name:
      data = L.Conv2D(
        filters=parameters['channels'],
        padding='same',
        kernel_size=parameters['kernel size'],
        activation=parameters['activation'],
        name='%s/conv-%d' % (name, i)
      )(data)
      continue

    if 'MLP Mixer' == Name:
      data = MixerConvLayer(
        token_mixing=parameters.get('token mixing', 512),
        channel_mixing=parameters.get('channel mixing', 512),
        name='%s/conv-mixer-%d' % (name, i)
      )(data)
      continue

    if 'Patches' == Name:
      data = Patches(
        patch_size=parameters['patch size'],
        name='%s/patches-%d' % (name, i)
      )(data)
      continue

    if 'CoordsGrid' == Name:
      parameters = {k: v for k, v in parameters.items() if k not in ['name']}
      parameters['name'] = '%s/coordsGrid-%d' % (name, i)
      data = CCoordsGridLayer(
        CCoordsEncodingLayer(
          N=parameters.get('N', 32),
          **parameters
        ),
        name='%s/coordsGrid-%d' % (name, i)
      )(data)
      continue

    if 'Transformer' == Name:
      parameters = {k: v for k, v in parameters.items()}      
      parameters['name'] = '%s/transformer-%d' % (name, i)
      parameters['intermediate_dim'] = parameters.pop('intermediate dim', 512)
      parameters['num_heads'] = parameters.pop('num heads', 8)
      data = TransformerBlock(**parameters)(data)
      continue

    if 'Reshape' == Name:
      shape = list(parameters['shape'])
      for j, sz in enumerate(shape):
        if sz <= -2:
          sz = data.shape[sz + 1]
        shape[j] = sz
        continue
      data = L.Reshape(
        shape,
        name='%s/reshape-%d' % (name, i)
      )(data)
      continue

    if 'MLP' == Name:
      parameters['name'] = '%s/mlp-%d' % (name, i)
      data = sMLP(**parameters)(data)
      continue
    
    raise NotImplementedError('Unknown layer: {}'.format(Name))
  return data

def _createGCMv2(dataShape, config, latentDim, name):
  data = L.Input(shape=dataShape)

  res = data
  for i, blockConfig in enumerate(config['downsample steps']):
    # downsample
    res = L.Conv2D(
      filters=blockConfig['channels'],
      kernel_size=blockConfig['kernel size'],
      strides=2,
      padding='same',
      activation='relu',
      name=name + '/downsample-%d' % (i + 1,)
    )(res)
    # convolutions
    for layerId in range(blockConfig['layers']):
      res = L.Conv2D(
        filters=blockConfig['channels'],
        kernel_size=blockConfig['kernel size'],
        padding='same',
        activation='relu',
        name=name + '/downsample-%d/layer-%d' % (i + 1, layerId + 1)
      )(res)
      continue
    continue

  return tf.keras.Model(inputs=[data], outputs=res, name=name)

def _createGlobalContextModel(X, config, latentDim, name):
  model = config.get('name', 'v1')
  if 'v1' == model: # simple convolutional model
    res = conv_block_from_config(
      data=X, config=config, defaults={
        'conv before': 0, # by default, no convolutions before the last layer
      }
    )
    # calculate global context
    latent = L.Flatten()(res)
    context = sMLP(sizes=config['mlp'], activation='relu', name=name + '/globalMixer')(latent)
    context = L.Dense(latentDim, activation=config['final activation'], name=name + '/dense-latent')(context)
    return context # end of 'v1' model
  
  if 'v2' == model:
    res = data = L.Input(shape=X.shape[1:])
    res = _createGCMv2(res.shape[1:], config, latentDim, name)(res)
    # calculate global context
    latent = L.Flatten()(res)
    context = sMLP(sizes=config['mlp'], activation='relu', name=name + '/globalMixer')(latent)
    context = L.Dense(latentDim, activation=config['final activation'], name=name + '/dense-latent')(context)
    model = tf.keras.Model(inputs=[data], outputs=context, name=name)
    return model(X) # end of 'v2' model
  
  raise NotImplementedError('Unknown global context model: {}'.format(model))

def _withPositionConfig(config, name):
  if config is None:
    print('[Encoder] Positions: No')
    return lambda x, _: x
  
  print('[Encoder] Positions: Yes')
  
  if isinstance(config, bool): config = { 'N': 32 }
  assert isinstance(config, dict), 'config must be a dictionary'

  def withPosition(x, i):
    if not config.get('stage-%d' % i, True): return x

    encoding = config.get('encoding', {})
    encoding = dict(**encoding)
    encoding['N'] = config.get('stage-%d N' % i, config.get('N', 32))
    return CCoordsGridLayer(
      CCoordsEncodingLayer(**encoding, name='%s/coordsGrid-%d/encoding' % (name, i)),
      name='%s/coordsGrid-%d' % (name, i)
    )(x)
  return withPosition

##################
def createEncoderHead_full(
  imgWidth,
  config,
  channels, downsampleSteps, latentDim, 
  ConvBeforeStage, ConvAfterStage, 
  localContext, globalContext,
  positionsConfigs,
  name
):
  assert config is not None, 'config must be a dictionary'
  assert isinstance(downsampleSteps, list) and (0 < len(downsampleSteps)), 'downsampleSteps must be a list of integers'
  data = L.Input(shape=(imgWidth, imgWidth, channels))
  
  withPosition = _withPositionConfig(positionsConfigs, name)
  res = data
  intermediate = []
  for i, sz in enumerate(downsampleSteps):
    if config.get('use downsampling', True):
      res = L.Conv2D(sz, 3, strides=2, padding='same', activation='relu')(res)
    res = withPosition(res, i) # add position encoding if needed
    for _ in range(ConvBeforeStage):
      res = L.Conv2D(sz, 3, padding='same', activation='relu')(res)

    # local context
    if not(localContext is None):
      intermediate.append(
        conv_block_from_config(
          data=res, config=localContext, defaults={
            'channels': sz,
            'channels last': latentDim, # last layer should have latentDim channels
          },
          name='%s/intermediate-%d' % (name, i)
        )
      )
    ################################
    for _ in range(ConvAfterStage):
      res = L.Conv2D(sz, 3, padding='same', activation='relu')(res)
    continue

  if not(globalContext is None): # global context
    res = withPosition(res, len(downsampleSteps))
    context = _createGlobalContextModel(res, globalContext, latentDim, name + '/globalContext')
  else: # no global context
    # return dummy context to keep the interface consistent
    context = L.Lambda(
      lambda x: tf.zeros((tf.shape(x)[0], 1), dtype=res.dtype)
    )(res)

  return tf.keras.Model(
    inputs=[data],
    outputs={
      'intermediate': intermediate, # intermediate representations
      'context': context, # global context
    },
    name=name
  )

class CEncoderHead(tf.keras.Model):
  def __init__(self, 
    config,
    downsampleSteps, latentDim, 
    ConvBeforeStage, ConvAfterStage, 
    localContext, globalContext,
    positionsConfigs,
    **kwargs
  ):
    super().__init__(**kwargs)
    self._config = config
    self._downsampleSteps = downsampleSteps
    self._latentDim = latentDim
    self._ConvBeforeStage = ConvBeforeStage
    self._ConvAfterStage = ConvAfterStage
    self._localContext = localContext
    self._globalContext = globalContext
    self._positionsConfigs = positionsConfigs
    return
  
  def build(self, inputShape):
    H, W, C = inputShape[1:]
    self._encoderHead = createEncoderHead_full(
      imgWidth=H, config=self._config,
      channels=C, downsampleSteps=self._downsampleSteps, latentDim=self._latentDim,
      ConvBeforeStage=self._ConvBeforeStage, ConvAfterStage=self._ConvAfterStage,
      localContext=self._localContext, globalContext=self._globalContext,
      positionsConfigs=self._positionsConfigs,
      name=self.name + '/EncoderHead'
    )
    self._encoderHead.build(inputShape)
    return super().build(inputShape)

  def call(self, src, training=None):
    return self._encoderHead(src, training=training)
'''
Simple encoder that takes image as input and returns corresponding latent vector with intermediate representations
'''
def createEncoderHead(
  config,
  downsampleSteps, latentDim, 
  ConvBeforeStage, ConvAfterStage, 
  localContext, globalContext,
  positionsConfigs,
  name
):
  return CEncoderHead(
    config=config,
    downsampleSteps=downsampleSteps,
    latentDim=latentDim,
    ConvBeforeStage=ConvBeforeStage,
    ConvAfterStage=ConvAfterStage,
    localContext=localContext,
    globalContext=globalContext,
    positionsConfigs=positionsConfigs,
    name=name
  )