update model

Author: Aaron Higuera

DUNE/RiceBayes/BNN_model.py

@@ -7,7 +7,7 @@
 tfd = tfp.distributions
 tfpl = tfp.layers
 
-num_samples = 1510865
+num_samples = 2023317
 
 
 def kl_approx(q, p, q_tensor):
@@ -184,7 +184,7 @@ def bayes_model(input_shape=(200,200,3)):
     x = layers.ReLU()(x)
     x = layers.MaxPooling2D(3, strides=2, padding='same')(x)
 
-    num_blocks = 3  # Increase the number of residual blocks
+    num_blocks = 4  # Increase the number of residual blocks
     filters = [32, 64, 128, 256, 512]   # Increase the number of filters in each block
 
     for i in range(num_blocks):

DUNE/RiceBayes/train.py

@@ -63,7 +63,8 @@ def on_epoch_end(self, epoch, logs=None):
 
         with open(self.file_path, 'w') as file:
             json.dump(self.history, file)
-   
+
+
 
 def nll(y_true, y_pred):
     """
@@ -110,8 +111,9 @@ def nll(y_true, y_pred):
 
     model.compile(optimizer=sgd_optimizer, loss=nll, metrics=['accuracy'])
 
-    checkpoint_filepath = '/Users/aaronhiguera/HEP/DUNE/Machine-Learning/DUNE/RiceBayes/'
-    checkpoint_callback = ModelCheckpoint(filepath=checkpoint_filepath, save_best_only=False)
+    checkpoint_filepath = '/home/higuera/RiceBNN/'
+    checkpoint_callback = ModelCheckpoint(filepath=checkpoint_filepath,
+                                          save_best_only=True, save_weights_only=True, monitor='val_loss')
 
     # Define the learning rate scheduler callback and history saver
     lr_scheduler = LearningRateSchedulerPlateau(factor=0.5, patience=5, min_lr=1e-6)
@@ -124,7 +126,7 @@ def nll(y_true, y_pred):
     validation_generator = DataGenerator(partition['validation'], **params)    
 
     model.fit(train_generator,validation_data=validation_generator,
-              epochs=args.num_epochs, callbacks=[lr_scheduler, history_saver, early_stopper])
+              epochs=args.num_epochs, callbacks=[lr_scheduler, history_saver, early_stopper, checkpoint_callback])
 
     # for inferences need to save weights
     weights = args.test_name+'.h5'