particle

networks/example_PCNN.py

+import torch
+import torch.nn as nn
+
+
+class ResNetUnit(nn.Module):
+    r"""Parameters
+    ----------
+    in_channels : int
+        Number of channels in the input vectors.
+    out_channels : int
+        Number of channels in the output vectors.
+    strides: tuple
+        Strides of the two convolutional layers, in the form of (stride0, stride1)
+    """
+
+    def __init__(self, in_channels, out_channels, strides=(1, 1), **kwargs):
+
+        super(ResNetUnit, self).__init__(**kwargs)
+        self.conv1 = nn.Conv1d(in_channels, out_channels, kernel_size=3, stride=strides[0], padding=1)
+        self.bn1 = nn.BatchNorm1d(out_channels)
+        self.conv2 = nn.Conv1d(out_channels, out_channels, kernel_size=3, stride=strides[1], padding=1)
+        self.bn2 = nn.BatchNorm1d(out_channels)
+        self.relu = nn.ReLU()
+        self.dim_match = True
+        if not in_channels == out_channels or not strides == (1, 1):  # dimensions not match
+            self.dim_match = False
+            self.conv_sc = nn.Conv1d(in_channels, out_channels, kernel_size=1,
+                                     stride=strides[0] * strides[1], bias=False)
+
+    def forward(self, x):
+        identity = x
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.conv2(x)
+        x = self.bn2(x)
+        x = self.relu(x)
+        # print('resnet unit', identity.shape, x.shape, self.dim_match)
+        if self.dim_match:
+            return identity + x
+        else:
+            return self.conv_sc(identity) + x
+
+
+class ResNet(nn.Module):
+    r"""Parameters
+    ----------
+    features_dims : int
+        Input feature dimensions.
+    num_classes : int
+        Number of output classes.
+    conv_params : list
+        List of the convolution layer parameters. 
+        The first element is a tuple of size 1, defining the transformed feature size for the initial feature convolution layer.
+        The following are tuples of feature size for multiple stages of the ResNet units. Each number defines an individual ResNet unit.
+    fc_params: list
+        List of fully connected layer parameters after all EdgeConv blocks, each element in the format of
+        (n_feat, drop_rate)
+    """
+
+    def __init__(self, features_dims, num_classes,
+                 conv_params=[(32,), (64, 64), (64, 64), (128, 128)],
+                 fc_params=[(512, 0.2)],
+                 for_inference=False,
+                 **kwargs):
+
+        super(ResNet, self).__init__(**kwargs)
+        self.conv_params = conv_params
+        self.num_stages = len(conv_params) - 1
+        self.fts_conv = nn.Sequential(
+            nn.BatchNorm1d(features_dims),
+            nn.Conv1d(
+                in_channels=features_dims, out_channels=conv_params[0][0],
+                kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm1d(conv_params[0][0]),
+            nn.ReLU())
+
+        # define ResNet units for each stage. Each unit is composed of a sequence of ResNetUnit block
+        self.resnet_units = nn.ModuleDict()
+        for i in range(self.num_stages):
+            # stack units[i] layers in this stage
+            unit_layers = []
+            for j in range(len(conv_params[i + 1])):
+                in_channels, out_channels = (conv_params[i][-1], conv_params[i + 1][0]) if j == 0 \
+                    else (conv_params[i + 1][j - 1], conv_params[i + 1][j])
+                strides = (2, 1) if (j == 0 and i > 0) else (1, 1)
+                unit_layers.append(ResNetUnit(in_channels, out_channels, strides))
+
+            self.resnet_units.add_module('resnet_unit_%d' % i, nn.Sequential(*unit_layers))
+
+        # define fully connected layers
+        fcs = []
+        for idx, layer_param in enumerate(fc_params):
+            channels, drop_rate = layer_param
+            in_chn = conv_params[-1][-1] if idx == 0 else fc_params[idx - 1][0]
+            fcs.append(nn.Sequential(nn.Linear(in_chn, channels), nn.ReLU(), nn.Dropout(drop_rate)))
+        fcs.append(nn.Linear(fc_params[-1][0], num_classes))
+        if for_inference:
+            fcs.append(nn.Softmax(dim=1))
+        self.fc = nn.Sequential(*fcs)
+
+    def forward(self, points, features, lorentz_vectors, mask):
+        # x: the feature vector, (N, C, P)
+        if mask is not None:
+            features = features * mask
+        x = self.fts_conv(features)
+        for i in range(self.num_stages):
+            x = self.resnet_units['resnet_unit_%d' % i](x)  # (N, C', P'), P'<P due to kernal_size>1 or stride>1
+
+        # global average pooling
+        x = x.mean(dim=-1)  # (N, C')
+        # fully connected
+        x = self.fc(x)  # (N, out_chn)
+        return x
+
+
+def get_model(data_config, **kwargs):
+    conv_params = [(32,), (64, 64), (64, 64), (128, 128)]
+    fc_params = [(512, 0.2)]
+
+    pf_features_dims = len(data_config.input_dicts['pf_features'])
+    num_classes = len(data_config.label_value)
+    model = ResNet(pf_features_dims, num_classes,
+                   conv_params=conv_params,
+                   fc_params=fc_params)
+
+    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()

networks/example_PFN.py

+import torch
+import torch.nn as nn
+
+
+class ParticleFlowNetwork(nn.Module):
+    r"""Parameters
+    ----------
+    input_dims : int
+        Input feature dimensions.
+    num_classes : int
+        Number of output classes.
+    layer_params : list
+        List of the feature size for each layer.
+    """
+
+    def __init__(self, input_dims, num_classes,
+                 Phi_sizes=(100, 100, 128),
+                 F_sizes=(100, 100, 100),
+                 use_bn=True,
+                 for_inference=False,
+                 **kwargs):
+
+        super(ParticleFlowNetwork, self).__init__(**kwargs)
+        # input bn
+        self.input_bn = nn.BatchNorm1d(input_dims) if use_bn else nn.Identity()
+        # per-particle functions
+        phi_layers = []
+        for i in range(len(Phi_sizes)):
+            phi_layers.append(nn.Sequential(
+                nn.Conv1d(input_dims if i == 0 else Phi_sizes[i - 1], Phi_sizes[i], kernel_size=1),
+                nn.BatchNorm1d(Phi_sizes[i]) if use_bn else nn.Identity(),
+                nn.ReLU())
+            )
+        self.phi = nn.Sequential(*phi_layers)
+        # global functions
+        f_layers = []
+        for i in range(len(F_sizes)):
+            f_layers.append(nn.Sequential(
+                nn.Linear(Phi_sizes[-1] if i == 0 else F_sizes[i - 1], F_sizes[i]),
+                nn.ReLU())
+            )
+        f_layers.append(nn.Linear(F_sizes[-1], num_classes))
+        if for_inference:
+            f_layers.append(nn.Softmax(dim=1))
+        self.fc = nn.Sequential(*f_layers)
+
+    def forward(self, points, features, lorentz_vectors, mask):
+        # x: the feature vector initally read from the data structure, in dimension (N, C, P)
+        x = self.input_bn(features)
+        x = self.phi(x)
+        if mask is not None:
+            x = x * mask.bool().float()
+        x = x.sum(-1)
+        return self.fc(x)
+
+
+def get_model(data_config, **kwargs):
+    Phi_sizes = (128, 128, 128)
+    F_sizes = (128, 128, 128)
+    input_dims = len(data_config.input_dicts['pf_features'])
+    num_classes = len(data_config.label_value)
+    model = ParticleFlowNetwork(input_dims, num_classes, Phi_sizes=Phi_sizes,
+                                F_sizes=F_sizes, use_bn=kwargs.get('use_bn', 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()

networks/example_ParticleNet.py

+import torch
+from weaver.nn.model.ParticleNet import ParticleNet
+
+'''
+Link to the full model implementation:
+https://github.com/hqucms/weaver-core/blob/main/weaver/nn/model/ParticleNet.py
+'''
+
+
+class ParticleNetWrapper(torch.nn.Module):
+    def __init__(self, **kwargs) -> None:
+        super().__init__()
+        self.mod = ParticleNet(**kwargs)
+
+    def forward(self, points, features, lorentz_vectors, mask):
+        return self.mod(points, features, mask)
+
+
+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()

networks/example_ParticleNet_finetune.py

+import torch
+import torch.nn as nn
+from weaver.nn.model.ParticleNet import ParticleNet
+
+'''
+Link to the full model implementation:
+https://github.com/hqucms/weaver-core/blob/main/weaver/nn/model/ParticleNet.py
+'''
+
+
+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()

networks/example_ParticleTransformer.py

+import torch
+from weaver.nn.model.ParticleTransformer import ParticleTransformer
+from weaver.utils.logger import _logger
+
+'''
+Link to the full model implementation:
+https://github.com/hqucms/weaver-core/blob/main/weaver/nn/model/ParticleTransformer.py
+'''
+
+
+class ParticleTransformerWrapper(torch.nn.Module):
+    def __init__(self, **kwargs) -> None:
+        super().__init__()
+        self.mod = ParticleTransformer(**kwargs)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'mod.cls_token', }
+
+    def forward(self, points, features, lorentz_vectors, mask):
+        return self.mod(features, v=lorentz_vectors, mask=mask)
+
+
+def get_model(data_config, **kwargs):
+
+    cfg = dict(
+        input_dim=len(data_config.input_dicts['pf_features']),
+        num_classes=len(data_config.label_value),
+        # network configurations
+        pair_input_dim=4,
+        use_pre_activation_pair=False,
+        embed_dims=[128, 512, 128],
+        pair_embed_dims=[64, 64, 64],
+        num_heads=8,
+        num_layers=8,
+        num_cls_layers=2,
+        block_params=None,
+        cls_block_params={'dropout': 0, 'attn_dropout': 0, 'activation_dropout': 0},
+        fc_params=[],
+        activation='gelu',
+        # misc
+        trim=True,
+        for_inference=False,
+    )
+    cfg.update(**kwargs)
+    _logger.info('Model config: %s' % str(cfg))
+
+    model = ParticleTransformerWrapper(**cfg)
+
+    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()

networks/example_ParticleTransformer_finetune.py

+import torch
+import torch.nn as nn
+from weaver.nn.model.ParticleTransformer import ParticleTransformer
+from weaver.utils.logger import _logger
+
+'''
+Link to the full model implementation:
+https://github.com/hqucms/weaver-core/blob/main/weaver/nn/model/ParticleTransformer.py
+'''
+
+
+class ParticleTransformerWrapper(nn.Module):
+    def __init__(self, **kwargs) -> None:
+        super().__init__()
+
+        in_dim = kwargs['embed_dims'][-1]
+        fc_params = kwargs.pop('fc_params')
+        num_classes = kwargs.pop('num_classes')
+        self.for_inference = kwargs['for_inference']
+
+        fcs = []
+        for out_dim, drop_rate in fc_params:
+            fcs.append(nn.Sequential(nn.Linear(in_dim, out_dim), nn.ReLU(), nn.Dropout(drop_rate)))
+            in_dim = out_dim
+        fcs.append(nn.Linear(in_dim, num_classes))
+        self.fc = nn.Sequential(*fcs)
+
+        kwargs['num_classes'] = None
+        kwargs['fc_params'] = None
+        self.mod = ParticleTransformer(**kwargs)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'mod.cls_token', }
+
+    def forward(self, points, features, lorentz_vectors, mask):
+        x_cls = self.mod(features, v=lorentz_vectors, mask=mask)
+        output = self.fc(x_cls)
+        if self.for_inference:
+            output = torch.softmax(output, dim=1)
+        return output
+
+
+def get_model(data_config, **kwargs):
+
+    cfg = dict(
+        input_dim=len(data_config.input_dicts['pf_features']),
+        num_classes=len(data_config.label_value),
+        # network configurations
+        pair_input_dim=4,
+        use_pre_activation_pair=False,
+        embed_dims=[128, 512, 128],
+        pair_embed_dims=[64, 64, 64],
+        num_heads=8,
+        num_layers=8,
+        num_cls_layers=2,
+        block_params=None,
+        cls_block_params={'dropout': 0, 'attn_dropout': 0, 'activation_dropout': 0},
+        fc_params=[],
+        activation='gelu',
+        # misc
+        trim=True,
+        for_inference=False,
+    )
+    cfg.update(**kwargs)
+    _logger.info('Model config: %s' % str(cfg))
+
+    model = ParticleTransformerWrapper(**cfg)
+
+    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()

test_QuarkGluon_demo.sh

+#!/bin/bash
+
+set -x
+
+source env.sh
+
+echo "args: $@"
+
+# set the dataset dir via `DATADIR_QuarkGluon`
+DATADIR=${DATADIR_QuarkGluon}
+[[ -z $DATADIR ]] && DATADIR='./datasets/QuarkGluon'
+
+# set a comment via `COMMENT`
+suffix=${COMMENT}
+
+model=$1
+extraopts=""
+modelopts="networks/example_ParticleTransformer.py"  # ParT
+
+# /public/home/mashun/gaonengsuo/particle_transformer/training/QuarkGluon/ParT/20241121-093112_example_ParticleTransformer_ranger_lr0.001_batch512/net_best_epoch_state.pt
+
+weaver --predict \
+    --data-test "${DATADIR}/test_file_*.parquet" \
+    --data-config data/QuarkGluon/qg_${FEATURE_TYPE}.yaml \
+    --network-config $modelopts \
+    --model-prefix  training/QuarkGluon/ParT/20241121-093112_example_ParticleTransformer_ranger_lr0.001_batch512/net_best_epoch_state.pt\
+    --batch-size 512 \
+    --gpus 0 \
+    --predict-output pred.root \

train_JetClass.sh

+#!/bin/bash
+
+set -x
+
+source env.sh
+
+echo "args: $@"
+
+# set the dataset dir via `DATADIR_JetClass`
+DATADIR=${DATADIR_JetClass}
+[[ -z $DATADIR ]] && DATADIR='./datasets/JetClass'
+
+# set a comment via `COMMENT`
+suffix=${COMMENT}
+
+# set the number of gpus for DDP training via `DDP_NGPUS`
+NGPUS=${DDP_NGPUS}
+[[ -z $NGPUS ]] && NGPUS=1
+if ((NGPUS > 1)); then
+    CMD="torchrun --standalone --nnodes=1 --nproc_per_node=$NGPUS -- $(which weaver) --backend nccl"
+else
+    CMD="weaver"
+fi
+
+epochs=1
+samples_per_epoch=$((10000 * 1024 / $NGPUS))
+samples_per_epoch_val=$((10000 * 128))
+dataopts="--num-workers 2 --fetch-step 0.01"
+
+# PN, PFN, PCNN, ParT
+model=$1
+if [[ "$model" == "ParT" ]]; then
+    modelopts="networks/example_ParticleTransformer.py --use-amp"
+    batchopts="--batch-size 512 --start-lr 1e-3"
+elif [[ "$model" == "PN" ]]; then
+    modelopts="networks/example_ParticleNet.py"
+    batchopts="--batch-size 512 --start-lr 1e-2"
+elif [[ "$model" == "PFN" ]]; then
+    modelopts="networks/example_PFN.py"
+    batchopts="--batch-size 4096 --start-lr 2e-2"
+elif [[ "$model" == "PCNN" ]]; then
+    modelopts="networks/example_PCNN.py"
+    batchopts="--batch-size 4096 --start-lr 2e-2"
+else
+    echo "Invalid model $model!"
+    exit 1
+fi
+
+# "kin", "kinpid", "full"
+FEATURE_TYPE=$2
+[[ -z ${FEATURE_TYPE} ]] && FEATURE_TYPE="full"
+
+if ! [[ "${FEATURE_TYPE}" =~ ^(full|kin|kinpid)$ ]]; then
+    echo "Invalid feature type ${FEATURE_TYPE}!"
+    exit 1
+fi
+
+# currently only Pythia
+SAMPLE_TYPE=Pythia
+
+
+$CMD \
+    --data-train \
+    "HToBB:${DATADIR}/${SAMPLE_TYPE}/train_100M/HToBB_*.root" \
+    "HToCC:${DATADIR}/${SAMPLE_TYPE}/train_100M/HToCC_*.root" \
+    "HToGG:${DATADIR}/${SAMPLE_TYPE}/train_100M/HToGG_*.root" \
+    "HToWW2Q1L:${DATADIR}/${SAMPLE_TYPE}/train_100M/HToWW2Q1L_*.root" \
+    "HToWW4Q:${DATADIR}/${SAMPLE_TYPE}/train_100M/HToWW4Q_*.root" \
+    "TTBar:${DATADIR}/${SAMPLE_TYPE}/train_100M/TTBar_*.root" \
+    "TTBarLep:${DATADIR}/${SAMPLE_TYPE}/train_100M/TTBarLep_*.root" \
+    "WToQQ:${DATADIR}/${SAMPLE_TYPE}/train_100M/WToQQ_*.root" \
+    "ZToQQ:${DATADIR}/${SAMPLE_TYPE}/train_100M/ZToQQ_*.root" \
+    "ZJetsToNuNu:${DATADIR}/${SAMPLE_TYPE}/train_100M/ZJetsToNuNu_*.root" \
+    --data-val "${DATADIR}/${SAMPLE_TYPE}/val_5M/*.root" \
+    --data-test \
+    "HToBB:${DATADIR}/${SAMPLE_TYPE}/test_20M/HToBB_*.root" \
+    "HToCC:${DATADIR}/${SAMPLE_TYPE}/test_20M/HToCC_*.root" \
+    "HToGG:${DATADIR}/${SAMPLE_TYPE}/test_20M/HToGG_*.root" \
+    "HToWW2Q1L:${DATADIR}/${SAMPLE_TYPE}/test_20M/HToWW2Q1L_*.root" \
+    "HToWW4Q:${DATADIR}/${SAMPLE_TYPE}/test_20M/HToWW4Q_*.root" \
+    "TTBar:${DATADIR}/${SAMPLE_TYPE}/test_20M/TTBar_*.root" \
+    "TTBarLep:${DATADIR}/${SAMPLE_TYPE}/test_20M/TTBarLep_*.root" \
+    "WToQQ:${DATADIR}/${SAMPLE_TYPE}/test_20M/WToQQ_*.root" \
+    "ZToQQ:${DATADIR}/${SAMPLE_TYPE}/test_20M/ZToQQ_*.root" \
+    "ZJetsToNuNu:${DATADIR}/${SAMPLE_TYPE}/test_20M/ZJetsToNuNu_*.root" \
+    --data-config data/JetClass/JetClass_${FEATURE_TYPE}.yaml --network-config $modelopts \
+    --model-prefix training/JetClass/${SAMPLE_TYPE}/${FEATURE_TYPE}/${model}/{auto}${suffix}/net \
+    $dataopts $batchopts \
+    --samples-per-epoch ${samples_per_epoch} --samples-per-epoch-val ${samples_per_epoch_val} --num-epochs $epochs --gpus 0 \
+    --optimizer ranger \
+    --log logs/JetClass_${SAMPLE_TYPE}_${FEATURE_TYPE}_${model}_{auto}${suffix}.log \
+    --predict-output pred.root \
+    --tensorboard JetClass_${SAMPLE_TYPE}_${FEATURE_TYPE}_${model}${suffix} \
+    "${@:3}"

train_QuarkGluon.sh

+#!/bin/bash
+
+set -x
+
+source env.sh
+
+echo "args: $@"
+
+# set the dataset dir via `DATADIR_QuarkGluon`
+DATADIR=${DATADIR_QuarkGluon}
+[[ -z $DATADIR ]] && DATADIR='./datasets/QuarkGluon'
+
+# set a comment via `COMMENT`
+suffix=${COMMENT}
+
+# PN, PFN, PCNN, ParT
+model=$1
+extraopts=""
+if [[ "$model" == "ParT" ]]; then
+    modelopts="networks/example_ParticleTransformer.py --use-amp --optimizer-option weight_decay 0.01"
+    lr="1e-3"
+elif [[ "$model" == "ParT-FineTune" ]]; then
+    modelopts="networks/example_ParticleTransformer_finetune.py --use-amp --optimizer-option weight_decay 0.01"
+    lr="1e-4"
+    extraopts="--optimizer-option lr_mult (\"fc.*\",50) --lr-scheduler none"
+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"
+    extraopts="--batch-size 4096"
+elif [[ "$model" == "PCNN" ]]; then
+    modelopts="networks/example_PCNN.py"
+    lr="2e-2"
+    extraopts="--batch-size 4096"
+else
+    echo "Invalid model $model!"
+    exit 1
+fi
+
+# "kin", "kinpid", "kinpidplus"
+FEATURE_TYPE=$2
+[[ -z ${FEATURE_TYPE} ]] && FEATURE_TYPE="kinpid"
+
+if [[ "${FEATURE_TYPE}" == "kin" ]]; then
+    pretrain_type="kin"
+elif [[ "${FEATURE_TYPE}" =~ ^(kinpid|kinpidplus)$ ]]; then
+    pretrain_type="kinpid"
+else
+    echo "Invalid feature type ${FEATURE_TYPE}!"
+    exit 1
+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" \
+    --data-test "${DATADIR}/test_file_*.parquet" \
+    --data-config data/QuarkGluon/qg_${FEATURE_TYPE}.yaml --network-config $modelopts \
+    --model-prefix training/QuarkGluon/${model}/{auto}${suffix}/net \
+    --num-workers 1 --fetch-step 1 --in-memory --train-val-split 0.8889 \
+    --batch-size 512 --samples-per-epoch 1600000 --samples-per-epoch-val 200000 --num-epochs 1 --gpus 0 \
+    --start-lr $lr --optimizer ranger --log logs/QuarkGluon_${model}_{auto}${suffix}.log --predict-output pred.root \
+    --tensorboard QuarkGluon_${FEATURE_TYPE}_${model}${suffix} \
+    ${extraopts} "${@:3}"

train_TopLandscape.sh

+#!/bin/bash
+
+set -x
+
+source env.sh
+
+echo "args: $@"
+
+# set the dataset dir via `DATADIR_TopLandscape`
+DATADIR=${DATADIR_TopLandscape}
+[[ -z $DATADIR ]] && DATADIR='./datasets/TopLandscape'
+# set a comment via `COMMENT`
+suffix=${COMMENT}
+
+# PN, PFN, PCNN, ParT
+model=$1
+extraopts=""
+if [[ "$model" == "ParT" ]]; then
+    modelopts="networks/example_ParticleTransformer.py --use-amp --optimizer-option weight_decay 0.01"
+    lr="1e-3"
+elif [[ "$model" == "ParT-FineTune" ]]; then
+    modelopts="networks/example_ParticleTransformer_finetune.py --use-amp --optimizer-option weight_decay 0.01"
+    lr="1e-4"
+    extraopts="--optimizer-option lr_mult (\"fc.*\",50) --lr-scheduler none --load-model-weights models/ParT_kin.pt"
+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"
+    extraopts="--batch-size 4096"
+elif [[ "$model" == "PCNN" ]]; then
+    modelopts="networks/example_PCNN.py"
+    lr="2e-2"
+    extraopts="--batch-size 4096"
+else
+    echo "Invalid model $model!"
+    exit 1
+fi
+
+# "kin"
+FEATURE_TYPE=$2
+[[ -z ${FEATURE_TYPE} ]] && FEATURE_TYPE="kin"
+if [[ "${FEATURE_TYPE}" != "kin" ]]; then
+    echo "Invalid feature type ${FEATURE_TYPE}!"
+    exit 1
+fi
+
+weaver \
+    --data-train "${DATADIR}/train_file.parquet" \
+    --data-val "${DATADIR}/val_file.parquet" \
+    --data-test "${DATADIR}/test_file.parquet" \
+    --data-config data/TopLandscape/top_${FEATURE_TYPE}.yaml --network-config $modelopts \
+    --model-prefix training/TopLandscape/${model}/{auto}${suffix}/net \
+    --num-workers 1 --fetch-step 1 --in-memory \
+    --batch-size 512 --samples-per-epoch $((2400 * 512)) --samples-per-epoch-val $((800 * 512)) --num-epochs 20 --gpus 0 \
+    --start-lr $lr --optimizer ranger --log logs/TopLandscape_${model}_{auto}${suffix}.log --predict-output pred.root \
+    --tensorboard TopLandscape_${FEATURE_TYPE}_${model}${suffix} \
+    ${extraopts} "${@:3}"