ParticleNet¶

class MultiLayerPerceptron(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,
                layer_params=(1024, 256, 256),
                **kwargs):

        super(MultiLayerPerceptron, self).__init__(**kwargs)
        channels = [input_dims] + list(layer_params) + [num_classes]
        layers = []
        for i in range(len(channels) - 1):
            layers.append(nn.Sequential(nn.Linear(channels[i], channels[i + 1]),
                                        nn.ReLU()))
        self.mlp = nn.Sequential(*layers)

    def forward(self, x):
        # x: the feature vector initally read from the data structure, in dimension (N, C, P)
        x = x.flatten(start_dim=1) # (N, L), where L = C * P
        return self.mlp(x)

def get_model(data_config, **kwargs):
    layer_params = (1024, 256, 256)
    _, pf_length, pf_features_dims = data_config.input_shapes['pf_features']
    input_dims = pf_length * pf_features_dims
    num_classes = len(data_config.label_value)
    model = MultiLayerPerceptron(input_dims, num_classes, layer_params=layer_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'}}},
        }

    print(model, model_info)
    return model, model_info


def get_loss(data_config, **kwargs):
    return torch.nn.CrossEntropyLoss()

selection:
### use `&`, `|`, `~` for logical operations on numpy arrays
### can use functions from `math`, `np` (numpy), and `awkward` in the expression

new_variables:
### [format] name: formula
### can use functions from `math`, `np` (numpy), and `awkward` in the expression
is_bkg: np.logical_not(is_signal_new)

preprocess:
### method: [manual, auto] - whether to use manually specified parameters for variable standardization
method: manual
### data_fraction: fraction of events to use when calculating the mean/scale for the standardization
data_fraction:

inputs:
pf_features:
    length: 100
    vars:
    ### [format 1]: var_name (no transformation)
    ### [format 2]: [var_name,
    ###              subtract_by(optional, default=None, no transf. if preprocess.method=manual, auto transf. if preprocess.method=auto),
    ###              multiply_by(optional, default=1),
    ###              clip_min(optional, default=-5),
    ###              clip_max(optional, default=5),
    ###              pad_value(optional, default=0)]
        - Part_Etarel
        - Part_Phirel
        - [Part_E_log, 2, 1]
        - [Part_P_log, 2, 1]

labels:
### type can be `simple`, `custom`
### [option 1] use `simple` for binary/multi-class classification, then `value` is a list of 0-1 labels
type: simple
value: [
    is_signal_new, is_bkg
    ]
### [option 2] otherwise use `custom` to define the label, then `value` is a map
# type: custom
# value:
    # target_mass: np.where(fj_isQCD, fj_genjet_sdmass, fj_gen_mass)

observers:
- origIdx
- idx
- Part_E_tot
- Part_PX_tot
- Part_PY_tot
- Part_PZ_tot
- Part_P_tot
- Part_Eta_tot
- Part_Phi_tot

# weights:
### [option 1] use precomputed weights stored in the input files
# use_precomputed_weights: true
# weight_branches: [weight, class_weight]
### [option 2] compute weights on-the-fly using reweighting histograms

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)],
                **kwargs):

        super(ResNet, self).__init__(**kwargs)
        self.conv_params = conv_params
        self.num_stages = len(conv_params) - 1
        self.fts_conv = nn.Sequential(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))
        self.fc = nn.Sequential(*fcs)

    def forward(self, x):
        # x: the feature vector, (N, C, P)
        x = self.fts_conv(x)
        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.sum(dim=-1) / x.shape[-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'}}},
        }

    print(model, model_info)
    print(data_config.input_shapes)
    return model, model_info


def get_loss(data_config, **kwargs):
    return torch.nn.CrossEntropyLoss()

class EdgeConvBlock(nn.Module):
    r"""EdgeConv layer.
    Introduced in "`Dynamic Graph CNN for Learning on Point Clouds
    <https://arxiv.org/pdf/1801.07829>`__".  Can be described as follows:
    .. math::
    x_i^{(l+1)} = \max_{j \in \mathcal{N}(i)} \mathrm{ReLU}(
    \Theta \cdot (x_j^{(l)} - x_i^{(l)}) + \Phi \cdot x_i^{(l)})
    where :math:`\mathcal{N}(i)` is the neighbor of :math:`i`.
    Parameters
    ----------
    in_feat : int
        Input feature size.
    out_feat : int
        Output feature size.
    batch_norm : bool
        Whether to include batch normalization on messages.
    """

    def __init__(self, k, in_feat, out_feats, batch_norm=True, activation=True, cpu_mode=False):
        super(EdgeConvBlock, self).__init__()
        self.k = k
        self.batch_norm = batch_norm
        self.activation = activation
        self.num_layers = len(out_feats)
        self.get_graph_feature = get_graph_feature_v2 if cpu_mode else get_graph_feature_v1

        self.convs = nn.ModuleList()
        for i in range(self.num_layers):
            self.convs.append(nn.Conv2d(2 * in_feat if i == 0 else out_feats[i - 1], out_feats[i], kernel_size=1, bias=False if self.batch_norm else True))

        if batch_norm:
            self.bns = nn.ModuleList()
            for i in range(self.num_layers):
                self.bns.append(nn.BatchNorm2d(out_feats[i]))

        if activation:
            self.acts = nn.ModuleList()
            for i in range(self.num_layers):
                self.acts.append(nn.ReLU())

        if in_feat == out_feats[-1]:
            self.sc = None
        else:
            self.sc = nn.Conv1d(in_feat, out_feats[-1], kernel_size=1, bias=False)
            self.sc_bn = nn.BatchNorm1d(out_feats[-1])

        if activation:
            self.sc_act = nn.ReLU()

    def forward(self, points, features):
        # points:   (N, C_pts, P)
        # features: (N, C_fts, P)
        # N: batch size, C: feature size per point, P: number of points

        topk_indices = knn(points, self.k) # (N, P, K)
        x = self.get_graph_feature(features, self.k, topk_indices) # (N, C_fts(0), P, K)

        for conv, bn, act in zip(self.convs, self.bns, self.acts):
            x = conv(x)  # (N, C', P, K)
            if bn:
                x = bn(x)
            if act:
                x = act(x)

        fts = x.mean(dim=-1)  # (N, C, P)

        # shortcut
        if self.sc:
            sc = self.sc(features)  # (N, C_out, P)
            sc = self.sc_bn(sc)
        else:
            sc = features

        return self.sc_act(sc + fts)  # (N, C_out, P)

class ParticleNet(nn.Module):
    r"""Parameters
    ----------
    input_dims : int
        Input feature dimensions (C_fts).
    num_classes : int
        Number of output classes.
    conv_params : list
        List of convolution parameters of EdgeConv blocks, each element in the format of (K, (C1, C2, C3)).
        K for the kNN number, C1,2,3 for convolution feature sizes of three layers in an EdgeConv block.
    fc_params: list
        List of fully connected layer parameters after all EdgeConv blocks, each element in the format of
        (n_feat, drop_rate)
    use_fusion: bool
        If true, concatenates all output features from each EdgeConv before the fully connected layer.
    use_fts_bn: bool
        If true, applies a batch norm before feeding to the EdgeConv block.
    use_counts: bool
        If true, uses the real count of points instead of the padded size (the max point size).
    for_inference: bool
        Whether this is an inference routine. If true, applies a softmax to the output.
    for_segmentation: bool
        Whether the model is set up for the point cloud segmentation (instead of classification) task. If true,
        does not merge the features after the last EdgeConv, and apply Conv1D instead of the linear layer.
        The output is hence each output_features per point, instead of output_features.
    """


    def __init__(self,
                input_dims,
                num_classes,
                conv_params=[(7, (32, 32, 32)), (7, (64, 64, 64))],
                fc_params=[(128, 0.1)],
                use_fusion=True,
                use_fts_bn=True,
                use_counts=True,
                for_inference=False,
                for_segmentation=False,
                **kwargs):
        super(ParticleNet, self).__init__(**kwargs)

        self.use_fts_bn = use_fts_bn
        if self.use_fts_bn:
            self.bn_fts = nn.BatchNorm1d(input_dims)

        self.use_counts = use_counts

        self.edge_convs = nn.ModuleList()
        for idx, layer_param in enumerate(conv_params):
            k, channels = layer_param
            in_feat = input_dims if idx == 0 else conv_params[idx - 1][1][-1]
            self.edge_convs.append(EdgeConvBlock(k=k, in_feat=in_feat, out_feats=channels, cpu_mode=for_inference))

        self.use_fusion = use_fusion
        if self.use_fusion:
            in_chn = sum(x[-1] for _, x in conv_params)
            out_chn = np.clip((in_chn // 128) * 128, 128, 1024)
            self.fusion_block = nn.Sequential(nn.Conv1d(in_chn, out_chn, kernel_size=1, bias=False), nn.BatchNorm1d(out_chn), nn.ReLU())

        self.for_segmentation = for_segmentation

        fcs = []
        for idx, layer_param in enumerate(fc_params):
            channels, drop_rate = layer_param
            if idx == 0:
                in_chn = out_chn if self.use_fusion else conv_params[-1][1][-1]
            else:
                in_chn = fc_params[idx - 1][0]
            if self.for_segmentation:
                fcs.append(nn.Sequential(nn.Conv1d(in_chn, channels, kernel_size=1, bias=False),
                                        nn.BatchNorm1d(channels), nn.ReLU(), nn.Dropout(drop_rate)))
            else:
                fcs.append(nn.Sequential(nn.Linear(in_chn, channels), nn.ReLU(), nn.Dropout(drop_rate)))
        if self.for_segmentation:
            fcs.append(nn.Conv1d(fc_params[-1][0], num_classes, kernel_size=1))
        else:
            fcs.append(nn.Linear(fc_params[-1][0], num_classes))
        self.fc = nn.Sequential(*fcs)

        self.for_inference = for_inference

    def forward(self, points, features, mask=None):
#         print('points:\n', points)
#         print('features:\n', features)
        if mask is None:
            mask = (features.abs().sum(dim=1, keepdim=True) != 0)  # (N, 1, P)
        points *= mask
        features *= mask
        coord_shift = (mask == 0) * 1e9
        if self.use_counts:
            counts = mask.float().sum(dim=-1)
            counts = torch.max(counts, torch.ones_like(counts))  # >=1

        if self.use_fts_bn:
            fts = self.bn_fts(features) * mask
        else:
            fts = features
        outputs = []
        for idx, conv in enumerate(self.edge_convs):
            pts = (points if idx == 0 else fts) + coord_shift
            fts = conv(pts, fts) * mask
            if self.use_fusion:
                outputs.append(fts)
        if self.use_fusion:
            fts = self.fusion_block(torch.cat(outputs, dim=1)) * mask

#         assert(((fts.abs().sum(dim=1, keepdim=True) != 0).float() - mask.float()).abs().sum().item() == 0)

        if self.for_segmentation:
            x = fts
        else:
            if self.use_counts:
                x = fts.sum(dim=-1) / counts  # divide by the real counts
            else:
                x = fts.mean(dim=-1)

        output = self.fc(x)
        if self.for_inference:
            output = torch.softmax(output, dim=1)
        # print('output:\n', output)
        return output

import torch
import torch.nn as nn
from utils.nn.model.ParticleNet import ParticleNet, FeatureConv


class ParticleNetTagger1Path(nn.Module):

    def __init__(self,
                pf_features_dims,
                num_classes,
                conv_params=[(7, (32, 32, 32)), (7, (64, 64, 64))],
                fc_params=[(128, 0.1)],
                use_fusion=True,
                use_fts_bn=True,
                use_counts=True,
                pf_input_dropout=None,
                for_inference=False,
                **kwargs):
        super(ParticleNetTagger1Path, self).__init__(**kwargs)
        self.pf_input_dropout = nn.Dropout(pf_input_dropout) if pf_input_dropout else None
        self.pf_conv = FeatureConv(pf_features_dims, 32)
        self.pn = ParticleNet(input_dims=32,
                            num_classes=num_classes,
                            conv_params=conv_params,
                            fc_params=fc_params,
                            use_fusion=use_fusion,
                            use_fts_bn=use_fts_bn,
                            use_counts=use_counts,
                            for_inference=for_inference)

    def forward(self, pf_points, pf_features, pf_mask):
        if self.pf_input_dropout:
            pf_mask = (self.pf_input_dropout(pf_mask) != 0).float()
            pf_points *= pf_mask
            pf_features *= pf_mask

        return self.pn(pf_points, self.pf_conv(pf_features * pf_mask) * pf_mask, pf_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)]
    use_fusion = True

    pf_features_dims = len(data_config.input_dicts['pf_features'])
    num_classes = len(data_config.label_value)
    model = ParticleNetTagger1Path(pf_features_dims, num_classes,
                            conv_params, fc_params,
                            use_fusion=use_fusion,
                            use_fts_bn=kwargs.get('use_fts_bn', False),
                            use_counts=kwargs.get('use_counts', True),
                            pf_input_dropout=kwargs.get('pf_input_dropout', None),
                            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'}}},
        }

    print(model, model_info)
    print(data_config.input_shapes)
    return model, model_info


def get_loss(data_config, **kwargs):
    return torch.nn.CrossEntropyLoss()

selection:
### use `&`, `|`, `~` for logical operations on numpy arrays
### can use functions from `math`, `np` (numpy), and `awkward` in the expression

new_variables:
### [format] name: formula
### can use functions from `math`, `np` (numpy), and `awkward` in the expression
pf_mask: awkward.JaggedArray.ones_like(Part_E)
is_bkg: np.logical_not(is_signal_new)

preprocess:
### method: [manual, auto] - whether to use manually specified parameters for variable standardization
method: manual
### data_fraction: fraction of events to use when calculating the mean/scale for the standardization
data_fraction:

inputs:
pf_points:
    length: 100
    vars:
        - Part_Etarel
        - Part_Phirel
pf_features:
    length: 100
    vars:
    ### [format 1]: var_name (no transformation)
    ### [format 2]: [var_name,
    ###              subtract_by(optional, default=None, no transf. if preprocess.method=manual, auto transf. if preprocess.method=auto),
    ###              multiply_by(optional, default=1),
    ###              clip_min(optional, default=-5),
    ###              clip_max(optional, default=5),
    ###              pad_value(optional, default=0)]
        - Part_Etarel
        - Part_Phirel
        - [Part_E_log, 2, 1]
        - [Part_P_log, 2, 1]
pf_mask:
    length: 100
    vars:
        - pf_mask

labels:
### type can be `simple`, `custom`
### [option 1] use `simple` for binary/multi-class classification, then `value` is a list of 0-1 labels
type: simple
value: [
    is_signal_new, is_bkg
    ]
### [option 2] otherwise use `custom` to define the label, then `value` is a map
# type: custom
# value:
    # target_mass: np.where(fj_isQCD, fj_genjet_sdmass, fj_gen_mass)

observers:
- origIdx
- idx
- Part_E_tot
- Part_PX_tot
- Part_PY_tot
- Part_PZ_tot
- Part_P_tot
- Part_Eta_tot
- Part_Phi_tot

# weights:
### [option 1] use precomputed weights stored in the input files
# use_precomputed_weights: true
# weight_branches: [weight, class_weight]
### [option 2] compute weights on-the-fly using reweighting histograms

#!/bin/bash

PREFIX=$1
MODEL_CONFIG=$2
DATA_CONFIG=$3
PATH_TO_SAMPLES=$4
WORKDIR=`pwd`

# Download miniconda
wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O miniconda_install.sh
bash miniconda_install.sh -b -p ${WORKDIR}/miniconda
export PATH=$WORKDIR/miniconda/bin:$PATH
pip install numpy pandas scikit-learn scipy matplotlib tqdm PyYAML
pip install uproot3 awkward0 lz4 xxhash
pip install tables
pip install onnxruntime-gpu
pip install tensorboard
pip install torch

# CUDA environment setup
export PATH=$PATH:/usr/local/cuda-10.2/bin
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda-10.2/lib64
export LIBRARY_PATH=$LIBRARY_PATH:/usr/local/cuda-10.2/lib64

# Clone weaver-benchmark
git clone --recursive https://github.com/colizz/weaver-benchmark.git
ln -s ../top_tagging weaver-benchmark/weaver/top_tagging
cd weaver-benchmark/weaver/
mkdir output

# Training, using 1 GPU
python train.py \
 --data-train ${PATH_TO_SAMPLES}'/prep/top_train_*.root' \
 --data-val ${PATH_TO_SAMPLES}'/prep/top_val_*.root' \
 --fetch-by-file --fetch-step 1 --num-workers 3 \
 --data-config top_tagging/data/${DATA_CONFIG} \
 --network-config top_tagging/networks/${MODEL_CONFIG} \
 --model-prefix output/${PREFIX} \
 --gpus 0 --batch-size 1024 --start-lr 5e-3 --num-epochs 20 --optimizer ranger \
 --log output/${PREFIX}.train.log

# Predicting score, using 1 GPU
python train.py --predict \
 --data-test ${PATH_TO_SAMPLES}'/prep/top_test_*.root' \
 --num-workers 3 \
 --data-config top_tagging/data/${DATA_CONFIG} \
 --network-config top_tagging/networks/${MODEL_CONFIG} \
 --model-prefix output/${PREFIX}_best_epoch_state.pt \
 --gpus 0 --batch-size 1024 \
 --predict-output output/${PREFIX}_predict.root

[ -d "runs/" ] && tar -caf output.tar output/ runs/ || tar -caf output.tar output/

Universe                = vanilla
executable              = run.sh
arguments               = <prefix> <model-config> <data-config> <your-path-to-samples>
#arguments              = mlp mlp_pf.py pf_features.yaml /eos/user/c/coli/public/weaver-benchmark/top_tagging/samples
output                  = job.$(ClusterId).$(ProcId).out
error                   = job.$(ClusterId).$(ProcId).err
log                     = job.$(ClusterId).log
should_transfer_files   = YES
when_to_transfer_output = ON_EXIT_OR_EVICT
transfer_output_files   = weaver-benchmark/weaver/output.tar
transfer_output_remaps  = "output.tar = output.$(ClusterId).$(ProcId).tar"
request_GPUs = 1
request_CPUs = 4
+MaxRuntime = 604800
queue