Cooperate paper revision: include the fine-tuning of ParticleNet

Author: colizz

README.md

@@ -70,22 +70,22 @@ Additional arguments will be passed directly to the `weaver` command, such as `-
 - using PyTorch's DistributedDataParallel:
 
 ```
-NGPUS=8 ./train_JetClass.sh ParT full --batch-size [batch_size_per_gpu] ...
+DDP_NGPUS=4 ./train_JetClass.sh ParT full --batch-size [batch_size_per_gpu] ...
 ```
 
 **To run the training on the QuarkGluon dataset:**
 
 ```
-./train_QuarkGluon.sh [ParT|ParT-FineTune|PN|PFN|PCNN] [kin|kinpid|kinpidplus] ...
+./train_QuarkGluon.sh [ParT|ParT-FineTune|PN|PN-FineTune|PFN|PCNN] [kin|kinpid|kinpidplus] ...
 ```
 
 **To run the training on the TopLandscape dataset:**
 
 ```
-./train_TopLandscape.sh [ParT|ParT-FineTune|PN|PFN|PCNN] [kin] ...
+./train_TopLandscape.sh [ParT|ParT-FineTune|PN|PN-FineTune|PFN|PCNN] [kin] ...
 ```
 
-The argument `ParT-FineTune` will run the fine-tuning using [models pre-trained on the JetClass dataset](models/).
+The argument `ParT-FineTune` or `PN-FineTune` will run the fine-tuning using [models pre-trained on the JetClass dataset](models/).
 
 ## Citations
 

models/ParticleNet_kin.pt

@@ -0,0 +1,61 @@
+import torch
+import torch.nn as nn
+from weaver.nn.model.ParticleNet import ParticleNet
+
+
+class ParticleNetWrapper(nn.Module):
+    def __init__(self, **kwargs) -> None:
+        super().__init__()
+
+        in_dim = kwargs['fc_params'][-1][0]
+        num_classes = kwargs['num_classes']
+        self.for_inference = kwargs['for_inference']
+
+        # finetune the last FC layer
+        self.fc_out = nn.Linear(in_dim, num_classes)
+
+        kwargs['for_inference'] = False
+        self.mod = ParticleNet(**kwargs)
+        self.mod.fc = self.mod.fc[:-1]
+
+    def forward(self, points, features, lorentz_vectors, mask):
+        x_cls = self.mod(points, features, mask)
+        output = self.fc_out(x_cls)
+        if self.for_inference:
+            output = torch.softmax(output, dim=1)
+        return output
+
+
+def get_model(data_config, **kwargs):
+    conv_params = [
+        (16, (64, 64, 64)),
+        (16, (128, 128, 128)),
+        (16, (256, 256, 256)),
+    ]
+    fc_params = [(256, 0.1)]
+
+    pf_features_dims = len(data_config.input_dicts['pf_features'])
+    num_classes = len(data_config.label_value)
+    model = ParticleNetWrapper(
+        input_dims=pf_features_dims,
+        num_classes=num_classes,
+        conv_params=kwargs.get('conv_params', conv_params),
+        fc_params=kwargs.get('fc_params', fc_params),
+        use_fusion=kwargs.get('use_fusion', False),
+        use_fts_bn=kwargs.get('use_fts_bn', True),
+        use_counts=kwargs.get('use_counts', True),
+        for_inference=kwargs.get('for_inference', False)
+    )
+
+    model_info = {
+        'input_names': list(data_config.input_names),
+        'input_shapes': {k: ((1,) + s[1:]) for k, s in data_config.input_shapes.items()},
+        'output_names': ['softmax'],
+        'dynamic_axes': {**{k: {0: 'N', 2: 'n_' + k.split('_')[0]} for k in data_config.input_names}, **{'softmax': {0: 'N'}}},
+    }
+
+    return model, model_info
+
+
+def get_loss(data_config, **kwargs):
+    return torch.nn.CrossEntropyLoss()

models/ParticleNet_kinpid.pt

@@ -39,8 +39,6 @@ def forward(self, points, features, lorentz_vectors, mask):
             output = torch.softmax(output, dim=1)
         return output
 
-        # return self.mod(features, v=lorentz_vectors, mask=mask)[:, [-2, 0]]
-
 
 def get_model(data_config, **kwargs):
 

networks/example_ParticleNet_finetune.py

@@ -26,6 +26,10 @@ elif [[ "$model" == "ParT-FineTune" ]]; then
 elif [[ "$model" == "PN" ]]; then
     modelopts="networks/example_ParticleNet.py"
     lr="1e-2"
+elif [[ "$model" == "PN-FineTune" ]]; then
+    modelopts="networks/example_ParticleNet_finetune.py"
+    lr="1e-3"
+    extraopts="--optimizer-option lr_mult (\"fc_out.*\",50) --lr-scheduler none"
 elif [[ "$model" == "PFN" ]]; then
     modelopts="networks/example_PFN.py"
     lr="2e-2"
@@ -55,6 +59,9 @@ fi
 if [[ "$model" == "ParT-FineTune" ]]; then
     modelopts+=" --load-model-weights models/ParT_${pretrain_type}.pt"
 fi
+if [[ "$model" == "PN-FineTune" ]]; then
+    modelopts+=" --load-model-weights models/ParticleNet_${pretrain_type}.pt"
+fi
 
 weaver \
     --data-train "${DATADIR}/train_file_*.parquet" \

networks/example_ParticleTransformer_finetune.py

@@ -25,6 +25,10 @@ elif [[ "$model" == "ParT-FineTune" ]]; then
 elif [[ "$model" == "PN" ]]; then
     modelopts="networks/example_ParticleNet.py"
     lr="1e-2"
+elif [[ "$model" == "PN-FineTune" ]]; then
+    modelopts="networks/example_ParticleNet_finetune.py"
+    lr="1e-3"
+    extraopts="--optimizer-option lr_mult (\"fc_out.*\",50) --lr-scheduler none --load-model-weights models/ParticleNet_kin.pt"
 elif [[ "$model" == "PFN" ]]; then
     modelopts="networks/example_PFN.py"
     lr="2e-2"