# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# ### use tf.keras subclass model to build your own custom model

# +
import argparse
import os
import numpy as np
import sys
from config import get_cfg_defaults

parser = argparse.ArgumentParser()
parser.add_argument(
    "--config-file",
    default=None,
    metavar="FILE",
    help="path to config file",
    type=str,
    )
parser.add_argument(
        "opts",
        help="Modify config options using the command-line",
        default=None,
        nargs=argparse.REMAINDER,
    )
args = parser.parse_args([])

# +
cfg = get_cfg_defaults()
if args.config_file is not None:
    cfg.merge_from_file(args.config_file)
if args.opts is not None:
    cfg.merge_from_list(args.opts)
cfg.freeze()

os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = str(cfg.SYSTEM.DEVICE)

# +
# model.py
import tensorflow as tf
import tensorflow.keras.layers as layers
import tensorflow.keras.models as models
import tensorflow.nn as F

class Conv_bn_relu(models.Model):
    """Stack blocks of Conv2D->BN->relu.
    
    Args:
      filters (int): numbers of filters of conv layer
      kernel_size (int): filter size
      strides (int): stride step
      data_format (str): channels_first or channels_last
      use_bias (bool): add bias to layer?
    Returns:
      tf.keras.model object
    """
    def __init__(self, filters, kernel_size=3, strides=1, data_format="channels_last",
                 use_bias=True, **kwargs):
        super(Conv_bn_relu, self).__init__(**kwargs)
        
        axis = -1 if data_format is "channels_last" else 1
        
        self.conv = Conv2DFixedPadding(filters=filters, kernel_size=kernel_size, strides=strides, use_bias=use_bias)
        self.normalize = layers.BatchNormalization(axis=axis)
        
    def call(self, x, training=True):
        x = self.conv(x)
        x = self.normalize(x, training=training)
        return F.relu(x)

class StackCNN(models.Model):
    """Stack all required layers together.
    
    Args:
      neurons_of_layers (list): list of filter size of convolution layers
      output_units (int): units of output node
    Returns:
      tf.keras.model object
    """
    def __init__(self, neurons_of_layers, output_units, **kwargs):
        super(StackCNN, self).__init__(**kwargs)
        
        self.layers_list = []
        for i, l in enumerate(neurons_of_layers):
            if (i+1) != len(neurons_of_layers):
                self.layers_list.append(Conv_bn_relu(filters=l, kernel_size=3, strides=1))
                self.layers_list.append(layers.MaxPooling2D(pool_size=(2,2)))
            else:
                self.layers_list.append(Conv_bn_relu(filters=l, kernel_size=3, strides=1))
                
        self.layers_list.append(layers.Flatten())
        self.layers_list.append(layers.Dense(units=output_units))
    
    def call(self, x, training=True):
        for l in self.layers_list:
            try:
                # some customized layer should give training flags
                x = l(x, training=training)
            except:
                # some original layers may not have training flags
                x = l(x)
        return F.softmax(x)

## Fixed Functions ##
def fixed_padding(inputs, kernel_size, data_format):
    """Pads the input along the spatial dimensions independently of input size.

    This function is copied from:
      https://github.com/tensorflow/models/blob/master/official/resnet/resnet_model.py

    Args:
      inputs: A tensor of size [batch, channels, height_in, width_in] or
        [batch, height_in, width_in, channels] depending on data_format.
      kernel_size: The kernel to be used in the conv2d or max_pool2d operation.
        Should be a positive integer.
      data_format: The input format ('channels_last' or 'channels_first').

    Returns:
      A tensor with the same format as the input with the data either intact
    (if kernel_size == 1) or padded (if kernel_size > 1).


    """

    pad_total = kernel_size - 1
    pad_beg = pad_total // 2
    pad_end = pad_total - pad_beg

    if data_format == 'channels_first':
        padded_inputs = tf.pad(tensor=inputs,
                               paddings=[[0, 0], [0, 0], [pad_beg, pad_end],
                                         [pad_beg, pad_end]])
    else:
        padded_inputs = tf.pad(tensor=inputs,
                               paddings=[[0, 0], [pad_beg, pad_end],
                                         [pad_beg, pad_end], [0, 0]])
    return padded_inputs


class Conv2DFixedPadding(models.Model):
    """Class for Strided 2-D convolution with explicit padding.

    The padding is consistent and is based only on `kernel_size`, not on the
    dimensions of `inputs` (as opposed to using `tf.layers.conv2d` alone).

    This class is based on:
      https://github.com/tensorflow/models/blob/master/official/resnet/resnet_model.py
    """

    def __init__(self, filters, kernel_size=3, strides=1, data_format="channels_last",
                 use_bias=True, **kwargs):
        super(Conv2DFixedPadding, self).__init__(**kwargs)
        self.kernel_size = kernel_size
        self.data_format = data_format
        self.strides = strides

        self.conv = layers.Conv2D(filters=filters, kernel_size=kernel_size,
                                  strides=strides, padding=('SAME' if strides == 1 else 'VALID'),
                                  use_bias=use_bias, data_format=data_format)
    
    def call(self, x):
        if self.strides > 1:
            x = fixed_padding(x, self.kernel_size, self.data_format)
        return self.conv(x)


# -

from tensorflow.keras.datasets.cifar10 import load_data

# +
train, valid = load_data()
x_train, y_train = train
x_valid, y_valid = valid

x_train = x_train / 255.
x_valid = x_valid / 255.
y_train = tf.keras.utils.to_categorical(y_train, 10)
y_valid = tf.keras.utils.to_categorical(y_valid, 10)
# -

input_layer = layers.Input(shape=[32,32,3])
module = StackCNN(neurons_of_layers=[32,32,64], output_units=10)(input_layer)
model = models.Model(inputs=input_layer, outputs=module)
model.compile(loss="categorical_crossentropy", optimizer=tf.train.AdamOptimizer(), metrics=["acc"])

model.summary()
model.get_layer('stack_cnn').summary()

model.fit(x_train, y_train, 
          epochs=cfg.TRAIN.EPOCHS, 
          batch_size=cfg.TRAIN.BATCH_SIZE, 
          shuffle=True, validation_data=(x_valid, y_valid))


# +
def build_example_model():
    input_layer = layers.Input(shape=(32,32,3))
    x = layers.Conv2D(filters=32, kernel_size=3, strides=1)(input_layer)
    x = layers.BatchNormalization()(x)
    x = layers.Activation("relu")(x)
    x = layers.MaxPooling2D(pool_size=(2,2))(x)
    
    x = layers.Conv2D(filters=32, kernel_size=3, strides=1)(x)
    x = layers.BatchNormalization()(x)
    x = layers.Activation("relu")(x)
    x = layers.MaxPooling2D(pool_size=(2,2))(x)
    
    x = layers.Conv2D(filters=64, kernel_size=3, strides=1)(x)
    x = layers.BatchNormalization()(x)
    x = layers.Activation("relu")(x)
    
    x = layers.GlobalAveragePooling2D()(x)
    x = layers.Dense(units=10, activation="softmax")(x)
    return models.Model(inputs=input_layer, outputs=x)

model = build_example_model()
model.compile(loss="categorical_crossentropy", optimizer=tf.train.AdamOptimizer(), metrics=["acc"])
model.summary()
# -

model.fit(x_train, y_train, 
          epochs=cfg.TRAIN.EPOCHS, 
          batch_size=cfg.TRAIN.BATCH_SIZE, 
          shuffle=True, validation_data=(x_valid, y_valid))

StarcoderdataJupyterScriptsDedupFiltered