README.md

particle_transformer

## 论文

`Particle Transformer for Jet Tagging`

* https://arxiv.org/abs/2202.03772

## 模型结构

该模型是一种基于 Transformer 的架构，通过成对粒子交互功能进行了增强，这些功能作为softmax之前的偏差合并到多头注意力中。

<img src="readme_imgs/arch.png" style="zoom:70%">

## 算法原理

该算法利用 Transformer 架构处理粒子云数据。Transformer 通过自注意力机制捕捉粒子之间的复杂关系，并学习粒子的全局特征。

<img src="readme_imgs/alg.png" style="zoom:100%">


## 环境配置

### Docker（方法一）
    
    docker pull image.sourcefind.cn:5000/dcu/admin/base/pytorch:2.1.0-ubuntu22.04-dtk24.04.3-py3.10

    docker run --shm-size 50g --network=host --name=pt --privileged --device=/dev/kfd --device=/dev/dri --group-add video --cap-add=SYS_PTRACE --security-opt seccomp=unconfined -v 项目地址(绝对路径):/home/ -v /opt/hyhal:/opt/hyhal:ro -it <your IMAGE ID> bash


### Dockerfile（方法二）

    docker build -t <IMAGE_NAME>:<TAG> .

    docker run --shm-size 50g --network=host --name=pt --privileged --device=/dev/kfd --device=/dev/dri --group-add video --cap-add=SYS_PTRACE --security-opt seccomp=unconfined -v 项目地址(绝对路径):/home/ -v /opt/hyhal:/opt/hyhal:ro -it <your IMAGE ID> bash


### Anaconda (方法三)

1、关于本项目DCU显卡所需的特殊深度学习库可从光合开发者社区下载安装：
https://developer.sourcefind.cn/tool/

    DTK驱动：dtk24.04.3
    python：python3.10
    torch: 2.1.0
    torchvision: 0.16.0

Tips：以上dtk驱动、python、torch等DCU相关工具版本需要严格一一对应

## 数据集

可使用脚本下载数据集:
```bash
./get_datasets.py [JetClass|QuarkGluon|TopLandscape] [-d DATA_DIR]
```

```
datasets/
├── JetClass
│   └── Pythia
│       ├── test_20M
│       ├── train_100M
│       └── val_5M
├── QuarkGluon
│   ├── test_file_*.parquet
│   └── train_file_*.parquet
└── TopLandscape
    ├── test_file.parquet
    ├── train_file.parquet
    └── val_file.parquet

```

## 训练

```bash
pip install 'weaver-core>=0.4'
```

```bash
# 在JetClass数据集上训练
./train_JetClass.sh [ParT|PN|PFN|PCNN] [kin|kinpid|full] ...

# 在QuarkGluon数据集上训练
./train_QuarkGluon.sh [ParT|ParT-FineTune|PN|PN-FineTune|PFN|PCNN] [kin|kinpid|kinpidplus] ...

# 在TopLandscape数据集上训练
./train_TopLandscape.sh [ParT|ParT-FineTune|PN|PN-FineTune|PFN|PCNN] [kin] ...
```

其中第一个参数表示使用的网络，分别是

- ParT: Particle Transformer
- PN: ParticleNet
- PFN: Particle Flow Network
- PCNN: P-CNN
- xxx-FineTune: 使用预训练模型微调

第二个参数表示输入的特征集合，分别是

- kin: only kinematic inputs
- kinpid: kinematic inputs + particle identification
- full: kinematic inputs + particle identification + trajectory displacement


### 多卡训练

```bash
# DP - pytorch
./train_JetClass.sh ParT full --gpus 0,1,2,3 --batch-size [total_batch_size] ...

# DDP - pytorch
DDP_NGPUS=4 ./train_JetClass.sh ParT full --batch-size [batch_size_per_gpu] ...
```

## 推理

```bash
bash test_QuarkGluon_demo.sh
```
注意：该推理脚本仅供参考，具体见[weaver-core](https://github.com/hqucms/weaver-core)。

## result

![alt text](readme_imgs/result.png)

### 精度

所有模型均选择ParT，以下结果仅展示不同加速卡在相同配置下的训练精度差异，并不代表可得到的最高精度。

仅记录AvgAcc.

|加速卡|JetClass|QuarkGluon(kinpidplus)|TopLandscape(kin)|
|:---:|:---:|:---:|:---:|
|K100_AI|0.6219|0.8495|0.93975|
|GPU|0.620|0.84921|0.93987|


## 应用场景

### 算法类别

`ai for science`

### 热点应用行业

`高能物理,医疗,金融`

## 源码仓库及问题反馈

* https://developer.sourcefind.cn/codes/modelzoo/particle_transformer_pytorch

## 参考资料

* https://github.com/jet-universe/particle_transformer
* https://github.com/hqucms/weaver-core