Tensorflow 2.0

bnorm_dropout.py

import os
import pdb
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3" 

import tensorflow as tf
if len(tf.config.list_physical_devices('GPU')):
    tf.config.experimental.set_memory_growth(tf.config.list_physical_devices('GPU')[0], True)

class ConvBlock3D(tf.keras.layers.Layer):

    def __init__(self, trainable):
        super(ConvBlock3D, self).__init__(name='ConvBlock3D')

        self.conv_layer = tf.keras.Sequential()
        self.conv_layer.add(tf.keras.layers.Conv3D(filters=10, kernel_size=3, trainable=trainable))
        self.conv_layer.add(tf.keras.layers.BatchNormalization(trainable=trainable))
    
    def call(self, x):

        return self.conv_layer(x)

class MyModel(tf.keras.Model):

    def __init__(self, trainable):
        super(MyModel, self).__init__(name='MyModel')

        self.convblock = ConvBlock3D(trainable)
    
    def call(self, x):

        return self.convblock(x)


if __name__ == "__main__":

    x = tf.random.normal((1,100,100,40,1))
    
    # Goal 1 - Check vars during training phase
    ## Expectation: running_mean, running_var must be trainable
    model = MyModel(trainable=True)
    _ = model(x, training=True)
    model.convblock.conv_layer.summary()
    print ('\n - model.trainable: ', model.trainable)
    for layer in model.layers:
        print (' - layer:{} || trainable: {}'.format(layer.name, layer.trainable))
        vars = layer.variables
        print (' -- var:{}  || trainable: {} ', vars[0].name, vars[0].trainable)
        vars_bn = [var for var in vars if 'normalization' in var.name]
        for var in vars_bn:
            print (' -- var:{}  || trainable: {} ', var.name, var.trainable)
    
    # Goal 2 - Check vars during inference phase
    model.trainable=False
    _ = model(x, training=False)
    model.convblock.conv_layer.summary()
    
    # Goal 3 - Check vars during inference phase (training=True for dropout at inference time)
    model.trainable=False
    _ = model(x, training=True)
    model.convblock.conv_layer.summary()
    print ('\n - model.trainable: ', model.trainable)
    for layer in model.layers:
        print (' - layer:{} || trainable: {}'.format(layer.name, layer.trainable))
        vars = layer.variables
        print (' -- var:{}  || trainable: {} ', vars[0].name, vars[0].trainable)
        vars_bn = [var for var in vars if 'normalization' in var.name]
        for var in vars_bn:
            print (' -- var:{}  || trainable: {} ', var.name, var.trainable)

gradienttape.py

# Ref: https://towardsdatascience.com/backpropagation-in-a-convolutional-layer-24c8d64d8509
## 1D conv with 2 channels

import os
import pdb
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"

import tensorflow as tf
if len(tf.config.list_physical_devices('GPU')):
    tf.config.experimental.set_memory_growth(tf.config.list_physical_devices('GPU')[0], True)


with tf.GradientTape(persistent=True, watch_accessed_variables=True) as tape:
    x = tf.ones([4,2]) # 2 channels with 4 inputs each (think of a timeseries with values for 4 timestamps)
    weights = tf.random.normal([2,2]) # 2 weights for each of the 2 channels
    mask = tf.constant([1.0, 0.0]) # masking to show what happens when one does not consider a particular channel (due to lack of data)
    tape.watch(x)
    tape.watch(weights)
    tape.watch(mask)

    y11 = tf.tensordot(x[0:2,0], weights[:,0], axes=1)
    y21 = tf.tensordot(x[1:3,0], weights[:,0], axes=1)
    y31 = tf.tensordot(x[2:4,0], weights[:,0], axes=1)
    y12 = tf.tensordot(x[0:2,1], weights[:,1], axes=1)
    y22 = tf.tensordot(x[1:3,1], weights[:,1], axes=1)
    y32 = tf.tensordot(x[2:4,1], weights[:,1], axes=1)

    y = tf.convert_to_tensor([[y11,y12],[y21,y22],[y31,y32]]) # shape=[3,2]
    ytrue = tf.ones(([3,2])) # random truth value
    loss1 = tf.math.reduce_sum((ytrue - y)*mask) # avg (full), shape=[1]
    loss2 = tf.math.reduce_sum((ytrue - y)*mask, axis=0) # avg (across labels only), shape = [1,2]
    # loss = loss * mask

print (' - loss1: ', loss1)
print (' - loss2: ', loss2)

gradients1 = tape.gradient(loss1, [weights])
gradients2 = tape.gradient(loss2, [weights])
print ('=================')
print (' - gradients1: ', gradients1)
print (' - gradients2: ', gradients2)
pdb.set_trace()

sample_dataloader.py

import os,sys
from pathlib import Path
MAIN_DIR = Path(__file__).parent.absolute().parent.absolute()
sys.path.append(str(MAIN_DIR))
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"

import pdb
import traceback
import numpy as np

import tensorflow as tf

class MICCAI2015Dataloader:
    """
    Loads .csv files for raw image and segmented masks
    :param path_raw_csv: The path of the .csv file containing paths to the raw images
    :param path_mask_csv: The path of the .csv file containing paths to the segmented mask images
    """

    def __init__(self, path_raw_csv, path_mask_csv):

        if Path(path_raw_csv).exists() and Path(path_mask_csv).exists():
            self.paths_raw = np.loadtxt(path_raw_csv, dtype='str')
            self.paths_mask = np.loadtxt(path_mask_csv, dtype='str')
        else:
            print (' - [ERROR][MICCAI2015Dataloader.__init__()] Path issues: ', Path(path_raw_csv).exists(), ' || ', Path(path_mask_csv).exists())
    
    def __len__(self):
        return len(self.paths_raw)
    
    def get_dataloader(self):
        try:
            dataset = tf.data.Dataset.from_generator(self.generator2D, (tf.float32, tf.float32), args=())
            dataset = dataset.map(self.aug_rotate, num_parallel_calls=tf.data.experimental.AUTOTUNE)
            return dataset

        except:
            traceback.print_exc()
            pdb.set_trace()

    def generator2D(self):
        try:
            idxs = np.arange(len(self.paths_raw))
            np.random.shuffle(idxs)
            for idx in idxs:
                slice_raw = np.load(self.paths_raw[idx])
                slice_mask = np.load(self.paths_mask[idx])

                if 1:
                    slice_raw = utils.crop_slice(slice_raw, config.H_START, config.H_END, config.W_START,config.W_END)
                    slice_mask = utils.crop_slice(slice_mask, config.H_START, config.H_END, config.W_START,config.W_END)

                if 1:
                    slice_raw = (slice_raw - np.mean(slice_raw)) / np.std(slice_raw)

                if 1:
                    label_ids = list(config.LABEL_MAP.values())
                    slice_mask_classes = np.zeros((slice_raw.shape[0], slice_raw.shape[1], len(label_ids)))
                    for label_id in label_ids:
                        label_idxs = np.argwhere(slice_mask == label_id)
                        slice_mask_classes[label_idxs[:,0], label_idxs[:,1], label_id-1] = 1

                yield (np.expand_dims(slice_raw, axis=2), slice_mask_classes)
        
        except:
            traceback.print_exc()
            pdb.set_trace()

    def generator2DCombinedMask(self):
        try:
            idxs = np.arange(len(self.paths_raw))
            np.random.shuffle(idxs)
            for idx in idxs:
                slice_raw = np.load(self.paths_raw[idx])
                slice_mask = np.load(self.paths_mask[idx])
                slice_raw = utils.crop_slice(slice_raw, config.H_START, config.H_END, config.W_START,config.W_END)
                slice_mask = utils.crop_slice(slice_mask, config.H_START, config.H_END, config.W_START,config.W_END)
                yield (np.expand_dims(slice_raw, axis=2), np.expand_dims(slice_mask,axis=2))

        except:
            traceback.print_exc()
            pdb.set_trace()

    def aug_rotate(self,x,y):
        try:
            if tf.random.uniform([], minval=0, maxval=1, dtype=tf.dtypes.float32) <= 0.5:
                rotate_count =  tf.random.uniform([], minval=1, maxval=4, dtype=tf.dtypes.int32)
                return tf.image.rot90(x, rotate_count), tf.image.rot90(y, rotate_count)
            else:
                return x,y
        except:
            traceback.print_exc()
            return x,y

if __name__ == "__main__":
    path_raw_csv  = Path(MAIN_DIR).joinpath('data','train',config.NAME_FOLDER_2D, config.NAME_FOLDER_2D_RAW, config.NAME_CSVFILE)
    path_mask_csv = Path(MAIN_DIR).joinpath('data','train',config.NAME_FOLDER_2D, config.NAME_FOLDER_2D_MASK, config.NAME_CSVFILE)

    for epoch in range(1):
        print (' ------------------------------------------- Epoch:{} -------------------------------------------'.format(epoch))
        dataset = MICCAI2015Dataloader(path_raw_csv, path_mask_csv)
        for idx, (X,Y) in enumerate(dataset.get_dataloader().batch(4)):
            print (X.shape, Y.shape)

sample_dataloader2.py

import os

os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3" # to avoid all those prints
os.environ["TF_GPU_THREAD_MODE"] = "gpu_private" # to avoid large "Kernel Launch Time"

import tensorflow as tf
if len(tf.config.list_physical_devices('GPU')):
    tf.config.experimental.set_memory_growth(tf.config.list_physical_devices('GPU')[0], True)

class Augmentations:

    def __init__(self):
        pass

    @tf.function
    def filter_even(self, x):
        if x % 2 == 0:
            return False
        else:
            return True

class Dataset:

    def __init__(self, aug, range_min=0, range_max=100):
        self.range_min = range_min
        self.range_max = range_max
        self.aug = aug

    def generator(self):
        dataset = tf.data.Dataset.from_generator(self._generator
                        , output_types=(tf.float32), args=())

        dataset = dataset.filter(self.aug.filter_even)

        return dataset
    
    def _generator(self):
        for item in range(self.range_min, self.range_max):
            yield(item)

class ZipDataset:

    def __init__(self, datasets):
        self.datasets = datasets
        self.datasets_generators = []
    
    def generator(self):
        for dataset in self.datasets:
            self.datasets_generators.append(dataset.generator())
        return tf.data.experimental.sample_from_datasets(self.datasets_generators)

if __name__ == "__main__":
    aug = Augmentations()
    dataset1 = Dataset(aug, 0, 100)
    dataset2 = Dataset(aug, 100, 200)
    dataset = ZipDataset([dataset1, dataset2])

    epochs = 2
    shuffle_buffer = 10
    batch_size = 4
    prefetch_buffer = 5

    dataset_gen = dataset.generator().shuffle(shuffle_buffer).batch(batch_size).prefetch(prefetch_buffer)
    # dataset_gen = dataset.generator().batch(batch_size).prefetch(prefetch_buffer) # this will output odd elements in sequence 

    for epoch in range(epochs):
        print ('\n ------------------ Epoch: {} ------------------'.format(epoch))
        for X in dataset_gen.repeat(1): # adding .repeat() in the loop allows you to easily control the end of the loop
            print (X)
        
        # Do some stuff at end of loop

sample_model_unet.py

import os,sys
from pathlib import Path
MAIN_DIR = Path(__file__).parent.absolute().parent.absolute()
sys.path.append(str(MAIN_DIR))
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"

tf.config.experimental.set_memory_growth(tf.config.list_physical_devices('GPU')[0], True)

import pdb
import traceback
import tensorflow as tf

class ConvBlock(tf.keras.layers.Layer):
    def __init__(self, filters, pool, kernel_size=(3,3), dropout=None, training=True, name=''):
        super(ConvBlock, self).__init__(name='{}_ConvBlock'.format(name))

        self.filters = filters
        self.pool = pool
        self.kernel_size = kernel_size
        self.dropout = dropout

        self.conv_layer = tf.keras.Sequential()
        for filter_id, filter_count in enumerate(self.filters):
            self.conv_layer.add(
                tf.keras.layers.Conv2D(filter_count, self.kernel_size, padding='same'
                        , kernel_regularizer=tf.keras.regularizers.l2(0.1)
                        , activation='relu'
                        , name='Conv_{}'.format(filter_id))
            )
            self.conv_layer.add(tf.keras.layers.BatchNormalization())
            if filter_id == 0 and self.dropout is not None:
                self.conv_layer.add(tf.keras.layers.Dropout(rate=self.dropout))
        
        self.pool_layer = tf.keras.layers.MaxPooling2D((2,2), strides=(2,2), name='Pool_{}'.format(self.name))
    
    def call(self, x):
        x = self.conv_layer(x)
        
        if self.pool is False:
            return x
        else:
            x_pool = self.pool_layer(x)
            return x, x_pool

class UpConvBlock(tf.keras.layers.Layer):
    def __init__(self, filters, kernel_size=(2,2), strides=(2, 2), name=''):
        super(UpConvBlock, self).__init__(name='{}_UpConv_Concat'.format(name))
        self.filters = filters
        self.kernel_size = kernel_size
        self.strides = strides

        self.upconv = tf.keras.layers.Conv2DTranspose(self.filters, self.kernel_size
                        , strides=self.strides
                        , kernel_regularizer=tf.keras.regularizers.l2(0.1)
                        , name='UpConv_{}'.format(self.name))
        
    def call(self, x1, x2):
        x1_upconv = self.upconv(x1)
        return tf.concat([x1_upconv, x2], axis=-1, name='Concat_{}'.format(self.name))

class ModelUNet(tf.keras.Model):
    def __init__(self, verbose=False):
        super(ModelUNet, self).__init__()

        self.class_count = len(config.LABEL_MAP)
        self.verbose = verbose

        self.convblock1 = ConvBlock(filters=[8, 8]    , pool=True , dropout=0.1, name='Block1')
        self.convblock2 = ConvBlock(filters=[16, 16]  , pool=True , dropout=0.1, name='Block2')
        self.convblock3 = ConvBlock(filters=[32, 32]  , pool=True , dropout=0.1, name='Block3')
        self.convblock4 = ConvBlock(filters=[64, 64]  , pool=True , dropout=0.2, name='Block4')
        self.convblock5 = ConvBlock(filters=[128, 128], pool=False, dropout=0.3, name='Block5')

        self.upconvblock6 = UpConvBlock(filters=64, name='Block6_1')
        self.convblock6   = ConvBlock(filters=[64,64], pool=False, dropout=0.2, name='Block6_2') 
        self.upconvblock7 = UpConvBlock(filters=32, name='Block7_1')
        self.convblock7   = ConvBlock(filters=[32,32], pool=False, dropout=0.1, name='Block7_2')
        self.upconvblock8 = UpConvBlock(filters=16, name='Block8_1')
        self.convblock8   = ConvBlock(filters=[16,16], pool=False, dropout=0.1, name='Block8_2')
        self.upconvblock9 = UpConvBlock(filters=8, name='Block9_1')
        self.convblock9   = ConvBlock(filters=[self.class_count,self.class_count], pool=False, dropout=0.1, name='Block9_2')
        self.convblock10  = tf.keras.layers.Conv2D(filters=9, kernel_size=(1,1), padding='same'
                            , activation='sigmoid' #activation=None
                            , name='Block10')

    def call(self,x):

        try:
            # conv1, pool1 = ConvBlock(filters=[8, 8]    , pool=True , dropout=0.1, name='Block1')(x)
            conv1, pool1 = self.convblock1(x)
            conv2, pool2 = self.convblock2(pool1)
            conv3, pool3 = self.convblock3(pool2)
            conv4, pool4 = self.convblock4(pool3)
            conv5        = self.convblock5(pool4)
            
            up6   = self.upconvblock6(conv5, conv4)
            conv6 = self.convblock6(up6)
            up7 = self.upconvblock7(conv6, conv3)
            conv7 = self.convblock7(up7)
            up8 = self.upconvblock8(conv7, conv2)
            conv8 = self.convblock8(up8)
            up9 = self.upconvblock9(conv8, conv1)
            conv9 = self.convblock9(up9)
            conv10 = self.convblock10(conv9)
            
            if self.verbose:
                print (' - x:', x.shape)
                print (' - conv1: ', conv1.shape, ' || pool1: ', pool1.shape)                
                print (' - conv2: ', conv2.shape, ' || pool2: ', pool2.shape)                
                print (' - conv3: ', conv3.shape, ' || pool3: ', pool3.shape)                
                print (' - conv4: ', conv4.shape, ' || pool4: ', pool4.shape)                
                print (' - conv5: ', conv5.shape)
                print (' - conv6: ', conv6.shape)                
                print (' - conv7: ', conv7.shape)                
                print (' - conv8: ', conv8.shape)                
                print (' - conv9: ', conv9.shape)
                print (' - conv10: ', conv10.shape)
                            
            return conv10

        except:
            traceback.print_exc()
            pdb.set_trace()
            
if __name__ == "__main__":
    X = tf.random.normal(shape=[1, 320, 320, 1])
    model = ModelUNet()
    y_predict = model(X, training=True)
    print (' - y_predict: ', y_predict.shape)
    print (' - model.variables: ', len(model.variables))
    print (' - model.trainable_variables: ', len(model.trainable_variables))
    for variable in model.variables: 
        print (' - ', variable.name)

segmentation_mask.py

import os
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"

import tensorflow as tf

@tf.function # Remove this to see the tf.print array values
def get_one_hot():
  label_ids = [0,5,10]
  mask_orig = tf.constant([[0,10], [0,10]], dtype=tf.float32) # [2,2]
  mask_onehot = tf.concat([tf.expand_dims(tf.math.equal(mask_orig, label_id),axis=-1) for label_id in label_ids], axis=-1) # [2,2,2]
  mask_label_present = tf.reduce_any(mask_onehot, axis=[0,1]) # [2]
  
  tf.print('\n - label_ids:{}'.format(label_ids))
  tf.print('\n - mask_orig:\n{}\n'.format(mask_orig))
  for id_, label_id in enumerate(label_ids):
      tf.print(' - mask_onehot:{}\n{}'.format(label_id, mask_onehot[:,:,id_]))
  tf.print('\n - mask_label_present:\n ', mask_label_present)

get_one_hot()

tensorflow_2.md

INSTALL TENSORFLOW

• Always install inside a conda env
• Download Anaconda
• Initialize your command prompt
  • For linux (conda init bash) OR Windows (conda init powershell)
• Create and activate env
  • conda create --name <env_name> python=3.7
  • conda activate <env_name>
• Check if you GPU is CUDA-enabled here

Method 1 - Using pip

• Install tensorflow
  • pip install tensorflow
• Install Cuda libraries (Ref)
  • conda install cudatoolkit (Search the correct CUDA/cuDNN version for each tensorflow version here)
    • Find the versions in conda using conda search cudatoolkit
  • conda install cudnn
    • Find the versions in conda using conda search cudnn
  • pip install tensorflow-probability Link
  • pip install tensorflow-addons Link
• To remove
  • pip uninstall tensorflow
  • conda remove cudatoolkit
  • conda remove cudnn

Using conda

• conda install -c anaconda tensorflow-gpu
• conda install tensorboard
• To remove
  • conda remove tensorflow-gpu

Check installation

• python -c "import tensorflow as tf;print('\n\n\n ====================== \n', tf.reduce_sum(tf.random.normal([1000, 1000])))"

• import tensorflow as tf; 
  tf.__version__; 
  tf.__file__; 
  tf.test.is_built_with_cuda()
  tf.config.list_physical_devices('CPU');
  tf.config.list_physical_devices('GPU');  
  

Packages to check

• conda env export --no-builds > tmp.yml
• Basic tensorflow (with GPU support) for version 2.3

  tensorflow==2.3.0
  tensorflow-estimator==2.3.0
  tensorboard==2.3.0
  keras-preprocessing==1.1.2
  cudatoolkit=10.1.243
  cudnn=7.6.5
  

CUDA and cuDNN

• import tensorflow as tf
  sys_details = tf.sysconfig.get_build_info()
  print (' - [TFlow Build Info] ver: ', tf.__version__, 'CUDA(major.minor):',  sys_details["cuda_version"], ' || cuDNN(major): ',       sys_details["cudnn_version"])
  

Other Refs

1. Nvidia Drivers
2. CUDA Driver from Nvidia installers

• non-conda method

2. nvidia-smi path: C:\Program Files\NVIDIA Corporation\NVSMI
3. tensorflow.python.platform.build_info.build_info
4. conda remove --name {myenv} --all

tf_gradient.py

"""
This script helps understand the differences in gradients when using different types of activations = [None, sigmoid, softmax]
"""

import os
import pdb
import traceback
import seaborn as sns
import matplotlib.pyplot as plt
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3" 

import tensorflow as tf
if len(tf.config.list_physical_devices('GPU')):
    tf.config.experimental.set_memory_growth(tf.config.list_physical_devices('GPU')[0], True)

def plot_hmaps_for_no_activation(show=False):

    with tf.GradientTape(persistent=True) as tape:
        y_predict = conv2D(x)
        loss = tf.math.reduce_sum(y_true - y_predict, axis=[1,2])
    
    dL_dY = tape.gradient(loss, y_predict)
    dY_dW = tape.gradient(y_predict, conv2D.variables)
    dL_dW = tape.gradient(loss, conv2D.variables)

    print (' - dL/dW = dL/dY . dY/dW: {} = {} . {}'.format(dL_dW[0].shape, dL_dY[0].shape, dY_dW[0].shape)) # (1, 1, 6, 6) = (50, 50, 6) * (1, 1, 6, 6)
    dL_dW_abs = tf.math.abs(dL_dW)
    sns.heatmap(dL_dW_abs[0,0,0,:,:], annot=True, vmin=0, vmax=0.02)
    plt.title('No activation (y=1, x=1e6)')
    if show:
        plt.show()
    else:
        plt.savefig('gradienthmaps_noactivation.png', bbox_inches='tight')

def plot_hmaps_for_sigmoid(show=False):

    f,axarr = plt.subplots(1,2, figsize=(15,10))

    with tf.GradientTape(persistent=True) as tape:
        z_predict = conv2D(x)
        y_predict = activation_sigmoid(z_predict)
        loss = tf.math.reduce_sum(y_true - y_predict, axis=[1,2])
        loss = loss * mask
    
    dL_dY = tape.gradient(loss, y_predict)
    dY_dZ = tape.gradient(y_predict, z_predict)
    dZ_dW = tape.gradient(z_predict, conv2D.variables)
    dL_dW = tape.gradient(loss, conv2D.variables)

    print (' - dL/dW = dL/dY . dY/dZ . dZ/dW: {} = {} . {} . {}'.format(dL_dW[0].shape, dL_dY[0].shape, dY_dZ[0].shape, dZ_dW[0].shape))
    dL_dW_abs = tf.math.abs(dL_dW)
    sns.heatmap(dL_dW_abs[0,0,0,:,:], ax=axarr[1], annot=True, vmin=0, vmax=0.02); axarr[1].set_title('Sigmoid activation with mask={}'.format(mask))
    
    with tf.GradientTape(persistent=True) as tape:
        z_predict = conv2D(x)
        y_predict = activation_sigmoid(z_predict)
        loss = tf.math.reduce_sum(y_true - y_predict, axis=[1,2])
    
    dL_dY = tape.gradient(loss, y_predict)
    dY_dZ = tape.gradient(y_predict, z_predict)
    dZ_dW = tape.gradient(z_predict, conv2D.variables)
    dL_dW = tape.gradient(loss, conv2D.variables)
    dL_dW_abs = tf.math.abs(dL_dW)
    sns.heatmap(dL_dW_abs[0,0,0,:,:], ax=axarr[0], annot=True, vmin=0, vmax=0.02); axarr[0].set_title('Sigmoid activation with no mask')
    
    if show:
        plt.show()
    else:
        plt.savefig('gradienthmaps_sigmoid.png', bbox_inches='tight')

def plot_hmaps_for_softmax(show=False):
    f,axarr = plt.subplots(1,2, figsize=(15,10))

    with tf.GradientTape(persistent=True) as tape:
        z_predict = conv2D(x)
        y_predict = activation_softmax(z_predict)
        loss = tf.math.reduce_sum(y_true - y_predict, axis=[1,2])
        loss = loss * mask
    
    dL_dY = tape.gradient(loss, y_predict)
    dY_dZ = tape.gradient(y_predict, z_predict)
    dZ_dW = tape.gradient(z_predict, conv2D.variables)
    dL_dW = tape.gradient(loss, conv2D.variables)

    print (' - dL/dW = dL/dY . dY/dZ . dZ/dW: {} = {} . {} . {}'.format(dL_dW[0].shape, dL_dY[0].shape, dY_dZ[0].shape, dZ_dW[0].shape))
    dL_dW_abs = tf.math.abs(dL_dW)
    sns.heatmap(dL_dW_abs[0,0,0,:,:], ax=axarr[1], annot=True, vmin=0, vmax=0.02); axarr[1].set_title('Softmax activation with mask={}'.format(mask))

    with tf.GradientTape(persistent=True) as tape:
        z_predict = conv2D(x)
        y_predict = activation_softmax(z_predict)
        loss = tf.math.reduce_sum(y_true - y_predict, axis=[1,2])
    
    dL_dY = tape.gradient(loss, y_predict)
    dY_dZ = tape.gradient(y_predict, z_predict)
    dZ_dW = tape.gradient(z_predict, conv2D.variables)
    dL_dW = tape.gradient(loss, conv2D.variables)
    dL_dW_abs = tf.math.abs(dL_dW)
    sns.heatmap(dL_dW_abs[0,0,0,:,:], ax=axarr[0], annot=True, vmin=0, vmax=0.02); axarr[0].set_title('Softmax activation with no mask')

    if show:
        plt.show()
    else:
        plt.savefig('gradienthmaps_softmax.png', bbox_inches='tight')

try:

    # Data
    # x      = tf.random.normal((1,50,50,6))
    x      = tf.fill((1,50,50,6), value=1e-6)
    y_true = tf.ones(((1,50,50,6)))
    mask   = tf.constant([1,1,1,1,0,0], dtype=tf.float32)

    # Model
    conv2D             = tf.keras.layers.Convolution2D(filters=6, kernel_size=(1,1), use_bias=False, activation=None)
    activation_sigmoid = tf.keras.layers.Activation(tf.nn.sigmoid)
    activation_softmax = tf.keras.layers.Activation(tf.nn.softmax)
    
    # Conv2D (w/o activation)
    if 1:    
        plot_hmaps_for_no_activation(show=False)
    
    # Conv2D (with sigmoid + no-mask/mask)
    elif 0:
        plot_hmaps_for_sigmoid(show=False)

    # Conv2D (with softmax + no-mask/mask)
    elif 1:
        plot_hmaps_for_softmax(show=False)

    pdb.set_trace()


except:
    traceback.print_exc()
    print ('\n\n ---------------- ERROR ----------------')
    pdb.set_trace()

weight_standardization.py

import os
import pdb
import traceback

os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3" 
import tensorflow as tf
if len(tf.config.list_physical_devices('GPU')):
    tf.config.experimental.set_memory_growth(tf.config.list_physical_devices('GPU')[0], True)
else:
    print (' - No GPU present!! Exiting ...')
    import sys; sys.exit(1)

class Conv3DWS(tf.keras.layers.Conv3D):
    """
    Ref
     - https://www.tensorflow.org/api_docs/python/tf/keras/layers/Conv3D
     - https://www.tensorflow.org/api_docs/python/tf/nn/conv3d
     - https://github.com/joe-siyuan-qiao/WeightStandardization
    """

    def __init__(self, filters, kernel_size=(3,3,3), strides=(1, 1, 1), padding='same'
                    , dilation_rate=(1,1,1)
                    , activation='relu'
                    , kernel_regularizer=None
                    , name=''):
        super(Conv3DWS, self).__init__(filters=filters, kernel_size=kernel_size, strides=strides, padding=padding
                        , dilation_rate=dilation_rate
                        , activation=activation
                        , kernel_regularizer=kernel_regularizer
                        , name=name)
    
    def call(self,x):

        # Step 1 - WS
        kernel_mean = tf.math.reduce_mean(self.kernel, axis=[0,1,2,3], keepdims=True, name='kernel_mean')
        kernel_std  = tf.math.reduce_std(self.kernel, axis=[0,1,2,3], keepdims=True, name='kernel_std')
        kernel_new  = (self.kernel - kernel_mean)/(kernel_std + tf.keras.backend.epsilon())
        
        # Step 2 - Convolution
        # output = tf.nn.conv3d(input=x, filters=kernel_new, strides=list((1,) + self.strides + (1,)), padding=self.padding.upper(), dilations=(1,) + self.dilation_rate + (1,)) # [Does not works due to padding=same]
        # output = tf.nn.conv3d(input=x, filters=kernel_new, strides=list((1,) + self.strides + (1,)), padding='VALID', dilations=(1,) + self.dilation_rate + (1,)) # [Works due to padding=valid]
        
        # Step 3 - Bias
        if self.use_bias:
            output = tf.nn.bias_add(output, self.bias, data_format=self._tf_data_format)
        
        # Step 4 - Activation and return
        if self.activation is not None:
            return self.activation(output)
        else:
            return output

if __name__ == "__main__":
    try:
        x = tf.random.normal((1,140,140,40,1))
        layers = Conv3DWS(filters=10, dilation_rate=(3,3,3))
        y = layers(x)
        print (' - y: ', y.shape)

    except:
        traceback.print_exc()
        pdb.set_trace()

A tip to get conda init to work for WIndows PowerShell

• Set-ExecutionPolicy -Scope CurrentUser Unrestricted -Force
• Ref