Future dev

README.md

@@ -2,8 +2,8 @@ HGCalML
 ===============================================================================
 
 Requirements
-  * DeepJetCore 3.X (``https://github.com/DL4Jets/DeepJetCore``)
-  * DeepJetCore 3.X container (or latest version in general)
+  * DeepJetCore 4.X (``https://github.com/DL4Jets/DeepJetCore``)
+  * DeepJetCore 4.X container (or latest version in general)
 
 For CERN (or any machine with cvmfs mounted), a script to start the latest container use this script:
 ```
@@ -23,7 +23,7 @@ sing=`which singularity`
 unset PATH
 cd
 
-$sing run -B /eos -B /afs $gpuopt /cvmfs/unpacked.cern.ch/registry.hub.docker.com/cernml4reco/deepjetcore3:latest
+$sing run -B /eos -B /afs $gpuopt /cvmfs/unpacked.cern.ch/registry.hub.docker.com/cernml4reco/deepjetcore4:latest
 ```
 
 The package follows the structure and logic of all DeepJetCore subpackages (also the example in DeepJetCore). So as a fresh starting point, it can be a good idea to follow the DeepJetCore example first.

Train/config_trainer.py

@@ -21,14 +21,14 @@
 from Layers import MixWhere
 from Layers import RaggedGravNet
 from Layers import PlotCoordinates
-from Layers import DistanceWeightedMessagePassing
+from Layers import DistanceWeightedMessagePassing, AccumulateNeighbours
 from Layers import LLFillSpace
 from Layers import LLExtendedObjectCondensation
-from Layers import DictModel,RaggedDictModel
+from Layers import DictModel
 from Layers import RaggedGlobalExchange
 from Layers import SphereActivation
 from Layers import Multi
-from Layers import ShiftDistance
+from Layers import ShiftDistance, KNN
 from Layers import LLRegulariseGravNetSpace
 from Regularizers import AverageDistanceRegularizer
 from model_blocks import tiny_pc_pool, condition_input
@@ -132,18 +132,21 @@ def config_model(Inputs, td, debug_outdir=None, plot_debug_every=2000):
     x = ScaledGooeyBatchNorm2(
             fluidity_decay=0.01,
             max_viscosity=0.999999,
+            learn=True,
             no_gaus=False)([x, is_track])
 
     x = ScaledGooeyBatchNorm2(
             fluidity_decay=0.01,
             max_viscosity=0.999999,
             invert_condition=True,
+            learn=True,
             no_gaus=False)([x, is_track])
 
     c_coords = prime_coords
     c_coords = ScaledGooeyBatchNorm2(
         name='batchnorm_ccoords',
         fluidity_decay=0.01,
+            learn=True,
         max_viscosity=0.999999)(c_coords)
     c_coords = PlotCoordinates(
         plot_every=plot_debug_every,
@@ -156,14 +159,20 @@ def config_model(Inputs, td, debug_outdir=None, plot_debug_every=2000):
     x = Concatenate()([x, c_coords, is_track])
     x = Dense(64, name='dense_pre_loop', activation=DENSE_ACTIVATION)(x) 
 
-    allfeat = []
+    allfeat = [c_coords]
     print("Available keys: ", pre_processed.keys())
 
+    ## testing ##
+
+    #nidx,dist = KNN(8, use_approximate_knn=True)([prime_coords, rs])
+    #x = Concatenate()([x, dist])
+    #x = Concatenate()([x,  DistanceWeightedMessagePassing([16])([x,nidx,dist]) ])
+
     ###########################################################################
     ### Loop over GravNet Layers ##############################################
     ###########################################################################
 
-    gravnet_regs = [0.01, 0.01, 0.01]
+    gravnet_reg = 0.01
 
     for i in range(GRAVNET_ITERATIONS):
 
@@ -189,14 +198,18 @@ def config_model(Inputs, td, debug_outdir=None, plot_debug_every=2000):
             )([x, rs])
 
         gndist = LLRegulariseGravNetSpace(
-                scale=gravnet_regs[i],
+                scale=gravnet_reg,
                 record_metrics=False,
                 name=f'regularise_gravnet_{i}')([gndist, prime_coords, gnnidx])
 
-        x_rand = random_sampling_block(
-                xgn, rs, gncoords, gnnidx, gndist, is_track,
-                reduction=6, layer_norm=True, name=f"RSU_{i}")
-        x_rand = ScaledGooeyBatchNorm2(**BATCHNORM_OPTIONS)(x_rand)
+        x = DistanceWeightedMessagePassing(
+                        [32,32,32,32,32,32,32],
+                         activation='elu')([x,gnnidx,gndist])
+
+        #x_rand = random_sampling_block(
+        #        xgn, rs, gncoords, gnnidx, gndist, is_track,
+        #        reduction=6, layer_norm=True, name=f"RSU_{i}")
+        #x_rand = ScaledGooeyBatchNorm2(**BATCHNORM_OPTIONS)(x_rand)
 
         gndist = AverageDistanceRegularizer(
             strength=1e-3,
@@ -214,8 +227,8 @@ def config_model(Inputs, td, debug_outdir=None, plot_debug_every=2000):
         # x_rand = ScalarMultiply(0.1)(x_rand)
         # gndist = ScalarMultiply(0.01)(gndist)
         # gncoords = ScalarMultiply(0.01)(gncoords)
-        x = Concatenate()([x_pre, xgn, x_rand, gndist, gncoords])
-        x = Dense(d_shape,
+        x = Concatenate()([x, x_pre, xgn, gndist, gncoords])
+        x = Dense(2*d_shape,
                   name=f"dense_post_gravnet_1_iteration_{i}",
                   activation=DENSE_ACTIVATION,
                   kernel_regularizer=DENSE_REGULARIZER)(x)
@@ -270,7 +283,7 @@ def config_model(Inputs, td, debug_outdir=None, plot_debug_every=2000):
 
     pred_beta = LLExtendedObjectCondensation(scale=1.,
                                              use_energy_weights=True,
-                                             record_metrics=False,
+                                             record_metrics=True,
                                              print_loss=True,
                                              name="ExtendedOCLoss",
                                              implementation = loss_implementation,
@@ -304,7 +317,7 @@ def config_model(Inputs, td, debug_outdir=None, plot_debug_every=2000):
         # 'no_noise_rs': pre_processed['no_noise_rs'],
         }
 
-    return RaggedDictModel(inputs=Inputs, outputs=model_outputs)
+    return DictModel(inputs=Inputs, outputs=model_outputs)
     #return DictModel(inputs=Inputs, outputs=model_outputs)