add prior weights from RiceRNN

Author: Aaron Higuera

DUNE/RiceBayes/BNN_model.py

@@ -1,3 +1,4 @@
+import numpy as np
 from functools import partial
 import tensorflow as tf
 from tensorflow.keras import layers, models
@@ -48,6 +49,21 @@ def prior(dtype, shape, name, trainable, add_variable_fn):
 
     return tfd.Independent(dist, reinterpreted_batch_ndims=batch_ndims)
 
+def customize_prior(dtype, shape, name, trainable, add_variable_fn):
+    """
+    Creates an customize normal distribution as a prior.
+
+    """
+    # Use weights from ResCNN as prior
+    prior_1stconv2d = np.load('/home/higuera/CNN/conv2d_weights_RiceRes.npy')
+    
+    mean = prior_1stconv2d  
+    
+    dist = tfd.MultivariateNormalDiag(loc = mean,
+                                        scale_diag = 1.5*tf.ones(shape))
+    batch_ndims = tf.size(dist.batch_shape_tensor())
+    
+    return tfd.Independent(dist, reinterpreted_batch_ndims=batch_ndims)
 
 def get_convolution_reparameterization(filters, kernel_size, activation, strides = 1,
                                         padding = 'SAME',
@@ -158,10 +174,12 @@ def bayes_model(input_shape=(200,200,3)):
     Returns:
         tf.keras.Model: The constructed Keras model.
     """
-    
+
     inputs = layers.Input(shape=input_shape, name='inputs')
 
-    x = get_convolution_reparameterization(16, 3, 'swish')(inputs)
+    x = get_convolution_reparameterization(32, 7, 'swish', strides=2, 
+                                           prior=customize_prior)(inputs)
+    
     x = layers.BatchNormalization(name='batchnorm_0')(x)
     x = layers.ReLU()(x)
     x = layers.MaxPooling2D(3, strides=2, padding='same')(x)

DUNE/RiceBayes/train.py

@@ -117,7 +117,7 @@ def nll(y_true, y_pred):
     lr_scheduler = LearningRateSchedulerPlateau(factor=0.5, patience=5, min_lr=1e-6)
     history_filename = args.test_name+'_training_history.json'
     history_saver = SaveHistoryToFile(history_filename)
-    early_stopper = EarlyStopping(monitor='val_loss', patience=3, min_delta=0.01, mode='min',
+    early_stopper = EarlyStopping(monitor='val_loss', patience=3, min_delta=0.001, mode='min',
                                   restore_best_weights=True)
 
     train_generator = DataGenerator(partition['train'], **params)