KL (Kullback Leibler) Divergence in Flipout Models (PyTorch + TFlow)

kl_tflow.py

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

import tensorflow as tf # v2.4 with CUDA11.0.0 & CuDNN8.0.0
import tensorflow_probability as tfp # v0.12.1

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

x = tf.random.normal((1,100,100,100,10)) # B, H,W,D,C
layer = tfp.layers.Convolution3DFlipout(filters=16, kernel_size=1) 
_ = layer(x)
print (' - [3D-big_ip] kl: ', layer.losses) # ~400

x = tf.random.normal((1,100,100,100,10))
layer = tfp.layers.Convolution3DFlipout(filters=16, kernel_size=3)
_ = layer(x)
print (' - [3D-big_ip] kl: ', layer.losses) # ~10k

x = tf.random.normal((1,100,100,100,32))
layer = tfp.layers.Convolution3DFlipout(filters=32, kernel_size=3)
_ = layer(x)
print (' - [3D-big_ip] kl: ', layer.losses) # ~70k

kl_torch.py

import torch # v1.8.0 with CUDA11.1 & CuDNN8.0.5
import bayesian_torch # https://github.com/IntelLabs/bayesian-torch
import bayesian_torch.layers

X     = torch.rand((1, 10, 100, 100, 100))  # B,C,H,W,D
layer = bayesian_torch.layers.flipout_layers.conv_flipout.Conv3dFlipout(in_channels=10, out_channels=16, kernel_size=1)
_, kl = layer(X) # kl~400

X     = torch.rand((1, 10, 100, 100, 100))
layer = bayesian_torch.layers.flipout_layers.conv_flipout.Conv3dFlipout(in_channels=10, out_channels=16, kernel_size=3)
_, kl = layer(X) # kl~10k

X     = torch.rand((1, 32, 100, 100, 100))
layer = bayesian_torch.layers.flipout_layers.conv_flipout.Conv3dFlipout(in_channels=32, out_channels=32, kernel_size=3)
_, kl = layer(X) # kl~70k