作者： He XingChen

最后更新： 2026-08-08

模型总览

SeisPolarity 提供一系列用于地震极性分类的深度学习模型。所有模型均使用 PyTorch 实现，并支持 GPU 加速。

可用模型

模型	输入长度	类别
Ross (SCSN)	400	3 (U/D/N)
Eqpolarity	600	2 (U/D)
DiTingMotion	128	3 (U/D/N)
CFM	160	2 (U/D)
RPNet	400	2 (U/D)
PolarCAP	64	2 (U/D)
APP	400	3 (U/D/N)

加载模型

预训练模型

从 Hugging Face 加载预训练模型：

from seispolarity.models import PPNet, RossNet, EqpolarityNet

# Ross 模型
model = RossNet(num_fm_classes=3)

# Eqpolarity 模型
model = EqpolarityNet()

# 通用 PPNet（用于 SCSN）
model = PPNet(num_fm_classes=3)

加载自定义权重

import torch
from seispolarity.models import PPNet

model = PPNet(num_fm_classes=3)

# 从检查点加载
checkpoint = torch.load("checkpoints/model.pth")
model.load_state_dict(checkpoint["model_state_dict"])

# 或直接加载
model.load_state_dict(torch.load("model_weights.pth"))

模型 API

SeisPolarity 中的所有模型都提供统一的 PyTorch 接口：

import torch
from seispolarity.models import PPNet

# 创建模型
model = PPNet(num_fm_classes=3)
model.eval()

# 准备输入
waveforms = torch.randn(10, 1, 400)  # (batch, channels, length)

# 前向传播
with torch.no_grad():
    logits = model(waveforms)
    predictions = logits.argmax(dim=1)  # 获取预测类别

模型结构

Ross (SCSN)

Ross 模型是针对 SCSN 数据优化的 CNN 结构。

输入：400 个采样点（100 Hz 采样率下 4 秒）

结构：

带批归一化的卷积层
最大池化
全连接层
用于正则化的 Dropout

from seispolarity.models import PPNet

model = PPNet(num_fm_classes=3)

Eqpolarity

Eqpolarity 是用于极性分类的深度 CNN 模型。

输入：600 个采样点（100 Hz 采样率下 6 秒）

from seispolarity.models import EqpolarityNet

model = EqpolarityNet()

DiTingMotion

基于运动的极性分类模型。

输入：128 个采样点（100 Hz 采样率下 1.28 秒）

from seispolarity.models import DiTingMotionNet

model = DiTingMotionNet()

CFM

用于极性检测的自定义结构。

输入：160 个采样点

from seispolarity.models import CFM

model = CFM()

RPNet

残差极性网络（Residual Polarity Network）。

输入：400 个采样点（100 Hz 采样率下 4 秒）

from seispolarity.models import RPNet

model = RPNet()

PolarCAP

用于极性分类的轻量级模型。

输入：64 个采样点（100 Hz 采样率下 0.64 秒）

from seispolarity.models import PolarCAP, PolarCAPLoss

model = PolarCAP()
loss_fn = PolarCAPLoss()

APP

自适应极性预测器（Adaptive Polarity Predictor）。

from seispolarity.models import APP

model = APP()

推理

批量推理

import torch
from seispolarity.models import PPNet
from seispolarity import WaveformDataset

# 加载模型
model = PPNet(num_fm_classes=3)
model.eval()
model.to("cuda")

# 加载数据集
dataset = WaveformDataset(path="data.hdf5", name="SCSN")
loader = dataset.get_dataloader(batch_size=1024)

# 推理
all_predictions = []
with torch.no_grad():
    for waveforms, _ in loader:
        waveforms = waveforms.to("cuda")
        logits = model(waveforms)
        predictions = logits.argmax(dim=1)
        all_predictions.append(predictions.cpu())

predictions = torch.cat(all_predictions)

单样本推理

import numpy as np
from seispolarity.models import PPNet
import torch

# 加载模型
model = PPNet(num_fm_classes=3)
model.eval()

# 单样本
waveform = np.random.randn(400)  # 单条波形
waveform = torch.FloatTensor(waveform).unsqueeze(0).unsqueeze(0)  # 添加 batch 和 channel 维度

# 推理
with torch.no_grad():
    logits = model(waveform)
    prediction = logits.argmax(dim=1).item()

# 解释结果
label_map = {0: "Up", 1: "Down", 2: "Unknown"}
print(f"Predicted polarity: {label_map[prediction]}")

模型输出

模型为每个类别输出 logits：

# 原始 logits
logits = model(waveforms)  # Shape: (batch_size, num_classes)

# 概率（softmax）
import torch.nn.functional as F
probabilities = F.softmax(logits, dim=1)

# 预测类别
predictions = logits.argmax(dim=1)

模型下载

预训练模型可在 Hugging Face 上获取：

from huggingface_hub import hf_hub_download
import torch

# 下载模型权重
model_path = hf_hub_download(
    repo_id="HeXingChen/SeisPolarity-Model",
    filename="ross_scsn.pth"
)

# 加载权重
state_dict = torch.load(model_path)
model.load_state_dict(state_dict)

自定义模型

创建自定义模型：

import torch.nn as nn
from seispolarity.models.base import BasePolarityModel

class CustomModel(BasePolarityModel):
    def __init__(self, num_fm_classes=3):
        super().__init__(num_fm_classes)
        # 在此处定义您的模型结构
        self.conv1 = nn.Conv1d(1, 32, kernel_size=5)
        self.fc = nn.Linear(32 * 396, num_fm_classes)

    def forward(self, x):
        # 前向传播
        x = self.conv1(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

详细的 API 文档请参见 API 参考。