# 地震学深度学习预训练模型

**Repository Path**: cangyeone/primedp

## Basic Information

- **Project Name**: 地震学深度学习预训练模型
- **Description**: 用于地震数据处理的多任务预训练模型。
- **Primary Language**: Unknown
- **License**: GPL-3.0
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 4
- **Forks**: 0
- **Created**: 2024-08-04
- **Last Updated**: 2025-01-24

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

### PRIME地震数据处理模型
我们的模型主要是基于Transformer的，使用Transformer来处理波形的特征。
![structure of PRIME-DP](fig/structure.png)

当前有几个模型
|模型名称|可训练参数数量|模型地址|状态|
|:-:|:-:|:-:|:-:|
|RNN版模型|77M|ckpt/primedp.rnn.pt|发布|
|Picker模型|0.5M|ckpt/primedp.picker.pt|发布|
|小模型|8.6M|ckpt/primedp.tinny.pt|发布|
|中模型|51M|ckpt/primedp.middle.pt|发布|
|大模型|1.3B|ckpt/primedp.large.pt|训练中|

### 1. 使用方式
```Python 
from prime.middle import PRIMEDP 
import torch 

model = PRIMEDP() 
#加载预训练模型
model.load_state_dict(torch.load("ckpt/primedp.middle.pt"))
#输入波形
x = torch.randn([32, 3, 10240])# N, C, T顺序
# 震相, 初动, 地震类型, 波形, 波形特征向量
phase, polar, event_type, wave, hidden = model(x) 
# 可以用于其他处理
# TO:比如使用hidden用于聚类分析
```

### 2. 迁移学习
以分类工作为例：
```Python 
# 定义可训练部分
for key, var in model.named_parameters():
    if var.dtype != torch.float32:continue # BN统计计数无梯度
    if "decoder_event_type" in key: # 仅有最后一层有out
        var.requires_grad = True
    else:
        var.requires_grad = False  
# 定义训练器    
optim = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), 1e-3, weight_decay=1e-1)
```
"ckpt\primedp.middel.classification.pt"用内蒙数据迁移训练的。

也可以使用特征来进行分类：
```Python 
import torch.nn as nn 

class Classification(nn.Module):
    def __init__(self) -> None:
        super().__init__() 
        self.decoder = nn.Linear(1024, 3) 
    def forward(self, x):
        y = self.decoder(x) 
        return y 
model = PRIMEDP() 
model.load_state_dict(torch.load("ckpt/primedp.middle.pt"))
x = torch.randn([32, 3, 10240])# N, C, T format. 
phase, polar, event_type, wave, hidden = model(x) 

decoder = Classification() 
vector = hidden[:, :, 0]#选择第0个输出，也可以使用其他输出
vector = vector.detach() 
y = decoder(vector) 
# TODO:定义损失函数并对decoder进行训练即可。
```

### 3. 更多使用方式
可以自行定义任务ID，来获取相应的任务。目前可以使用tinny，middle来完成。
任务解码器可以定义多个，即单一模型可以完成更多其他任务。训练解码器的时候需要给定task_id，目前01分别用于分类和单台定位，最多可以定义32个ID，即同一模型可以完成32种任务。但是需要注意的是我们的模型未必是训练充分的，即仅训练解码器参数可能精度不高，还需要将将更多模型参数加入训练。
```Python
model = PRIMEDP() 
model.load_state_dict(torch.load("ckpt/primedp.middle.pt"))
# 这里可以定义多任务ID，目前0是分类，1是单台定位，2以后自行定义
task_id = torch.Tensor([0, 1, 2], dtype=torch.long) 
x = torch.randn([32, 3, 10240])# N, C, T format. 
vect_task, vect_wave = model(x, task_id) 
vect_task=...#可以用于其他工作，当前shape为[32, 3]，3为task_id长度
```

如果任务ID选择更多长度比如32个。可以使用vect_wave来插值成波形长度用于去噪、缺失值填补等工作。


### 4. 联系方式
有任何使用问题可以联系：
蔡育埼，caiyuqiming@foxmail.com


### 5. 许可证
[GPLv3](LICENSE)