From 85b3ed5beff79cb35c9ca9efeb09e99b465790e1 Mon Sep 17 00:00:00 2001 From: cangyeone <305601470@qq.com> Date: Wed, 17 Dec 2025 16:48:16 +0800 Subject: [PATCH] Clarify ONNX export steps in README --- README.md | 132 ++++--- jit_picker_base.py | 88 +++++ makejit.eqt.py | 78 +--- makejit.eqtransformer.py | 725 +--------------------------------- makejit.pnsn.diff.py | 92 +---- makejit.pnsn.py | 80 +--- makejit.rnn.py | 76 +--- makejit.unet.py | 90 +---- makejit.unetpp.py | 90 +---- makeonnx.eqt.py | 61 +-- makeonnx.eqtransformer.py | 706 +-------------------------------- makeonnx.lppn.py | 100 +++-- makeonnx.pnsn.py | 59 +-- makeonnx.rnn.py | 59 +-- makeonnx.skynet.multiphase.py | 97 +---- makeonnx.skynet.py | 97 +---- makeonnx.unet.py | 70 +--- makeonnx.unetpp.py | 71 +--- onnx_picker_base.py | 74 ++++ 19 files changed, 503 insertions(+), 2342 deletions(-) create mode 100644 jit_picker_base.py create mode 100644 onnx_picker_base.py diff --git a/README.md b/README.md index d98c59d..1daa940 100755 --- a/README.md +++ b/README.md @@ -1,12 +1,11 @@ ### 1. Instructions for using the national 100Hz model -1. Training data, all models are based on the 2009-2019 national seismic network training of the 100Hz model, which can be directly used for continuous data picking. -2. The model training data is based on stations within a distance of 800km from the epicenter and includes PS wave data. -3. Currently, it has been tested based on three phases of ChinArray data, and the recall rate of RNN model manually labeled data is not less than 80%. -4. Different models' accuracy and speed are shown in the figure. -![](pickers/speed.jpg) - -#### The open sourced models -The open sourced models are as follows: +All models in this repository are trained on 2009-2019 national seismic network data at 100 Hz. They can be applied directly to continuous three-component waveforms for automatic phase picking. + +* Training covers stations within 800 km of the epicenter and includes P/S phases. +* PhaseNet, RNN and LPPN style models have been validated on ChinArray data with RNN recall ≥ 80% on manually labelled sets. +* Accuracy and speed comparisons are shown in `pickers/speed.jpg`. + +#### 1.1 Open sourced models |Model|Size(MB)|P-F1Score|Instrument|Sampling rate|Channel|Max distance|Range|Output phases| |:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:| |BRNN|1.9|0.857|Broad band|100Hz|EHZ|300km|Global|Pg、Sg| @@ -16,18 +15,18 @@ The open sourced models are as follows: |LPPN(Medium)|0.4|0.808|Broad band|100Hz|EHZ|300km|Global|Pg、Sg| |LPPN(Tinny)|0.3|0.757|Broad band|100Hz|EHZ|300km|Global|Pg、Sg| |UNet++|12|0.798|Broad band|100Hz|EHZ|300km|Global|Pg、Sg| -|pnsn(**The model is used in our article**)|1.9|0.781|Broad band, MEMS, |100Hz|EHZ|2000km|Global|Pg、Sg、Pn、Sn| -|pnsn.diff(**The model is used in our article**)|1.9|0.781|Broad band, MEMS, |100Hz|EHZ|2000km|Global|Pg、Sg、Pn、Sn| +|pnsn (**used in the paper**)|1.9|0.781|Broad band, MEMS|100Hz|EHZ|2000km|Global|Pg、Sg、Pn、Sn| +|pnsn.diff (**used in the paper**)|1.9|0.781|Broad band, MEMS|100Hz|EHZ|2000km|Global|Pg、Sg、Pn、Sn| |tele|1.9|0.800|Broad band|20Hz|EHZ|>3000km|Global|P| -|BRNN|1.9|0.807|Broad band|100Hz|Any|300km|Global|Pg、Sg| -pnsn.jit is the model used in our article for the first inference strategy and pnsn.diff.jit is the model used in our article for the second inference strategy. You can put any length of the three components data to the model. +`pickers/pnsn.jit` implements the first inference strategy and `pickers/pnsn.diff.jit` the second; both accept waveforms of arbitrary length. -Example usage: +#### 1.2 TorchScript quick start +The TorchScript models in `pickers/` ship with all post-processing (thresholding and non-maximum suppression) baked into the graph. They expect three-component waveforms resampled to 100 Hz and output `[phase_type, relative_sample, confidence]` for each pick. ```python import numpy as np # Import NumPy for numerical operations import torch # Import PyTorch for loading and running the model -import obspy # Import ObsPy for reading seismic waveform data (install via `pip install obspy`) +import obspy # Import ObsPy for reading seismic waveform data mname = "pickers/pnsn.jit" # Path to the TorchScript seismic phase picking model device = torch.device("cpu") # Set the inference device to CPU @@ -69,13 +68,13 @@ plt.show() # Display the plot -#### 1.1 Recommended models: -1. If accuracy is more important, pnsn can be used. We have tested it on mobile networks, dense networks, and fixed networks at the global level. -2. If memory is limited and speed is more important, LPPNM can be used. -3. If recall rate is low, we recommend using a threshold of 0.1 (pickers/rnn.01.jit), or using the PnSn model. Although the F1 score was low in testing, this was due to testing with manually labeled data within 2000km. -4. For some tasks that require confidence scores for each sampling point, an onnx model can be used. +#### 1.3 Recommended models +1. If accuracy is most important, prefer the pnsn/pnsn.diff variants (tested on mobile, dense, and fixed global networks). +2. If memory is tight or speed matters, choose LPPN models. +3. For low recall scenarios, lower the confidence threshold to 0.1 (for example `pickers/rnn.01.jit`) or use the Pn/Sn-aware models. +4. When per-sample confidence traces are required, use an ONNX model and handle post-processing externally. -#### 1.2 Pn and Sn phase picking model +#### 1.4 Pn and Sn phase picking model 1. In order to make the model more universal, we trained a new model using 2000km of manually labeled data. 2. The model is called rnn.pnsn.jit. 3. Based on the RNN model, it can simultaneously pick P, S, Pn, and Sn phases. @@ -85,7 +84,7 @@ plt.show() # Display the plot 7. The data needs to be sampled at 100Hz. 8.The accuracy has not been fully tested yet; only 10,000 waveforms of 102.4 seconds within 2000km from year 2020 were used for testing with results shown in the figure. 9.We found that after high-pass filtering (differentiation), the picking effect for large earthquakes was better; therefore we created a model for picking original + differentiated data as an example: makejit.pnsn.diff.py.Output models are: rnn.origdiff.pnsn.jit -10. The pnsn model was originally on 2022, however, we have updated it pnsn.jit and pnsn.diff.jit on 2025 in our article. The orignal accuracy is shown in [pickers/china.pnsn.jpg]. +10. The pnsn model was originally released in 2022, and updated `pnsn.jit` and `pnsn.diff.jit` were released in 2025 in our paper. The original accuracy is shown in [pickers/china.pnsn.jpg]. Call in python interface @@ -99,64 +98,83 @@ with torch.no_grad(): phase = y.cpu().numpy()# [Number of phases, 1P, 2S, 3Pn, 4Sn] ``` -#### 1.3 Distant Earthquake Picking Model -We have added a new model tele.rnn.jit for distant earthquake picking, which is used for picking the PS phase of distant earthquakes. +#### 1.5 Distant Earthquake Picking Model +We provide `tele.rnn.jit` for distant event picking. It outputs distant P/S phases at 20 Hz. + +#### 1.6 Environment and data prerequisites +The examples in this repository rely on common scientific Python packages: `torch`, `numpy`, `obspy`, `scipy`, `matplotlib`, and `tqdm` (see the imports in `picker.py`). The picker utilities assume three-component waveforms sampled at 100 Hz with channel names such as `BHE/BHN/BHZ` and file extensions ending in `.mseed` by default (see `config/picker.py`). ### 2. Model Usage Instructions We provide three types of model files: -1. .pt files in the ckpt folder, which can be used for transfer learning and easily transferred to local data. It is recommended to fix some trainable parameters during transfer training. -2. Models for picking any length are located in the pickers folder. - - .jit for direct use with PyTorch, which can directly output phase relative arrival time and phase type information. - - .onnx for use with onnxruntime library, which is lighter than PyTorch and suitable for picking on edge devices. Due to the simple API provided, post-processing needs to be done externally. -- The output format of .jit files is: [number of phases, phase type + relative arrival time + confidence], all jit files are like this. Phase types: 1:P, 2:S. -- The .onnx output has two parts: a probability prob and a time time; for example prob[i] represents the probability of different phase types at point i, it is a vector of length 3; time[i] represents the relative moment at point i. Time and prob need to be used together in order to perform picking. -- Example usage of .jit can be found in picker.jit.py -- Example usage of .onnx can be found in picker.onnx.py +1. `.pt` files in the `ckpt` folder, which can be used for transfer learning. Freeze some parameters when adapting to local data. +2. Models for picking any length are located in the `pickers` folder. + - `.jit` for direct use with PyTorch; post-processing is embedded in the graph and outputs `[phase_type, relative_sample, confidence]` per pick. + - `.onnx` for use with `onnxruntime`, suitable for edge devices. Use the `post` functions in `picker.onnx.py` or `picker.py` to apply the probability threshold (`a`) and non-maximum suppression window (`b`) to the raw `prob` and `time` outputs. +- `.jit` output format: `[number of phases, phase type + relative arrival time + confidence]`. Phase types: 1:P, 2:S (Pn/Sn models extend this list). +- `.onnx` outputs two tensors: `prob[i]` (per-sample class probabilities, length 3) and `time[i]` (relative sample index). Combine them with post-processing to form picks. +- Example usage of .jit can be found in `picker.jit.py`. +- Example usage of .onnx can be found in `picker.onnx.py`. #### 2.1 Using C Language Version Onnx Model -Due to the complexity of writing programs in C language, we have merged the time and prob outputs from onnx into a .merge.onnx version model where the vector format becomes: -[ [time length, number of categories,-,-], - [number of categories, noise probability,P-wave probability,S-wave probability], - [sample points, noise probability,P-wave probability,S-wave probability], - .....] -For examples using C language version programs please contact yuziye@cea-igp.ac.cn. +For C users, `.merge.onnx` files combine the `time` and `prob` outputs into a single array: +``` +[ [time length, number of categories, -, -], + [number of categories, noise probability, P-wave probability, S-wave probability], + [sample points, noise probability, P-wave probability, S-wave probability], + ... ] +``` +For example programs in C, contact yuziye@cea-igp.ac.cn. ### 3. make onnx and jit files -See the example programs makeonnx.xxx.jit and makejit.xxx.jit. In the .jit file: -```python -time_sel = torch.masked_select(ot, pc>0.3) -score = torch.masked_select(pc, pc>0.3) -``` -Here, 0.3 is the minimum confidence level, which seems reasonable at present. If you want to pick up more phases (and consequently more errors), you can lower this value. +#### 3.1 Building `.jit` pickers +All TorchScript pickers share the same interface via `jit_picker_base.py::SlidingWindowPicker`. Each `makejit.XXX.py` file simply: +1. Constructs the underlying network with `self.model = UNet()`/`BRNN()`/`EQTransformer()`, etc. +2. Loads a checkpoint whose keys are prefixed with `model.` (legacy checkpoints are also accepted and will be auto-prefixed). +3. Wraps the network with sliding-window preprocessing, softmax, and non-maximum suppression. +4. Saves the scripted model into `pickers/*.jit`. + +To rebuild the packaged TorchScript files, run the corresponding script (for example `python makejit.unet.py`, `python makejit.unetpp.py`, `python makejit.rnn.py`, `python makejit.pnsn.py`, or `python makejit.eqt.py`). The output `.jit` files include post-processing, so they return `[phase_type, relative_sample, confidence]` directly when you call `torch.jit.load`. + +Key thresholds baked into the picker interface: ```python -selidx = torch.masked_select(selidx, torch.abs(ref-ntime)>1000) -nprob = torch.masked_select(nprob, torch.abs(ref-ntime)>1000) -ntime = torch.masked_select(ntime, torch.abs(ref-ntime)>1000) +time_sel = torch.masked_select(ot, pc > 0.3) # confidence threshold +selidx = torch.masked_select(selidx, torch.abs(ref - ntime) > 1000) # NMS window (samples) ``` -Here, 1000 represents 1000 sampling points and signifies that only the phase with the highest probability within a window of length 1000 is picked up for the same type of phase. If it is believed that there may be multiple phases within a 10-second window, this value can be lowered. -**The onnx model can use config/picker.py for post-processing as it is outside of the model itself** - +* `0.3` is the default minimum confidence. Lower it to pick more candidates at the cost of extra false triggers. +* `1000` samples (10 seconds at 100 Hz) enforce a single pick per class within that window. Reduce the window if multiple phases are expected in short succession. + +#### 3.2 Building `.onnx` pickers +All ONNX pickers share the `OnnxSlidingWindowPicker` interface defined in `onnx_picker_base.py`. To regenerate the exported ONNX +files: +1. Run the corresponding script (for example `python makeonnx.unet.py`, `python makeonnx.unetpp.py`, `python makeonnx.rnn.py`, + `python makeonnx.pnsn.py`, or `python makeonnx.eqt.py`). +2. Each script builds the model (`self.model = UNet()`/`BRNN()`/`EQTransformer()`, etc.), loads checkpoints (auto-prefixing with + `model.` when needed), and wraps it with the shared sliding-window preprocessing. +3. Post-processing (probability threshold and NMS) remains outside the ONNX graph; reuse `config/picker.py` together with the + `post` helpers in `picker.onnx.py` or `picker.py` when running inference. ### 4. Directly picking up continuous data #### 4.1 Phase picking Phase picking provides a more convenient way to directly traverse the directory and pick up all phases. -```bash +```bash python picker.py -i path/to/data -o outputname -m pickers/rnn.jit -d device ``` 1. output file name.txt containing all picked phases 2. output file name.log containing processed data information -3. output file name.err containing problematic data information +3. output file name.err containing problematic data information The format of the output file is: ```text #path/to/file -phase name,relative time(s),confident,aboulute time(%Y-%m-%d %H:%M:%S.%f),SNR,AMP,station name,other information +phase name,relative time(s),confident,aboulute time(%Y-%m-%d %H:%M:%S.%f),SNR,AMP,station name,other information ``` +`picker.py` exposes the `-i/--input`, `-o/--output`, `-m/--model`, and `-d/--device` arguments (see `if __name__ == "__main__"` in the script) and uses the defaults from `config/picker.py` for details such as channel count (`nchannel=3`), sampling rate (`samplerate=100`), probability threshold for ONNX models (`prob=0.3`), and non-maximum suppression window (`nmslen=1000`). + #### 4.2 Seimic assosication -The goal of seismic association is to determine the number, location, and timing information of earthquakes from the phase picking results. Currently, there are 3 association algorithms provided: +The goal of seismic association is to determine the number, location, and timing information of earthquakes from the phase picking results. Currently, there are 3 association algorithms provided: 1. REAL methods [reallinker.py] 2. LPPN methods [fastlinker.py] 3. GaMMA methods [gammalinker.py] @@ -186,4 +204,10 @@ PHASE_PICKED_TIME_LAT_LON_TYPE_PROB_STATION_DIST_DELTA_ERROR# ``` ### Open Source License -GPLv3 \ No newline at end of file +GPLv3 + +### Related publication +* **Journal:** Journal of Geophysical Research: Machine Learning and Computation (Open Access) +* **Title:** *A Deep Learning Framework for Pg/Sg/Pn/Sn Phase Picking and Its Nationwide Implementation in Chinese Mainland* +* **DOI:** 10.1029/2025JH000944 +* **Status:** In Production diff --git a/jit_picker_base.py b/jit_picker_base.py new file mode 100644 index 0000000..93a9a85 --- /dev/null +++ b/jit_picker_base.py @@ -0,0 +1,88 @@ +import torch +import torch.nn as nn + + +class SlidingWindowPicker(nn.Module): + """ + Wrapper to standardize JIT picker interfaces. + + Each picker holds its underlying network in ``self.model``. The checkpoint + is expected to use keys prefixed with ``model.``; legacy checkpoints without + the prefix are supported by automatically adding it during load. + """ + + def __init__(self, model_ctor, ckpt_path=None, *, state_dict=None, + seqlen=6144, overlap=256, threshold=0.3, min_gap=1000): + super().__init__() + self.model = model_ctor() + self.n_stride = 1 + self.seqlen = seqlen + self.batchstride = seqlen - overlap + self.threshold = threshold + self.min_gap = min_gap + + if state_dict is None: + if ckpt_path is None: + raise ValueError("Either ckpt_path or state_dict must be provided") + state_dict = torch.load(ckpt_path, map_location="cpu") + if not any(k.startswith("model.") for k in state_dict.keys()): + state_dict = {f"model.{k}": v for k, v in state_dict.items()} + self.load_state_dict(state_dict, strict=False) + + def forward(self, x): + device = x.device + with torch.no_grad(): + T, _ = x.shape + batchlen = torch.ceil(torch.tensor(T / self.batchstride).to(device)) + idx = ( + torch.arange(0, self.seqlen, 1, device=device).unsqueeze(0) + + torch.arange(0, batchlen, 1, device=device).unsqueeze(1) + * self.batchstride + ) + idx = idx.clamp(min=0, max=T - 1).long() + wave = x.to(device)[idx, :] + wave = wave.permute(0, 2, 1) + wave -= torch.mean(wave, dim=2, keepdim=True) + maxv, _ = torch.max(torch.abs(wave), dim=2, keepdim=True) + wave /= (maxv + 1e-6) + + logits = self.model(wave) + if logits.dim() == 4: + logits = logits.squeeze(dim=3) + if logits.shape[1] > 1: + logits = logits.softmax(dim=1) + + B, C, T = logits.shape + tgrid = ( + torch.arange(0, T, 1, device=device).unsqueeze(0) * self.n_stride + + torch.arange(0, batchlen, 1, device=device).unsqueeze(1) * self.batchstride + ) + oc = logits.permute(0, 2, 1).reshape(-1, C) + ot = tgrid.squeeze().reshape(-1) + + output = [] + for itr in range(2): + pc = oc[:, itr + 1] + time_sel = torch.masked_select(ot, pc > self.threshold) + score = torch.masked_select(pc, pc > self.threshold) + _, order = score.sort(0, descending=True) + ntime = time_sel[order] + nprob = score[order] + select = -torch.ones_like(order) + selidx = torch.arange(0, order.numel(), 1, dtype=torch.long, device=device) + while True: + if nprob.numel() < 1: + break + ref = ntime[0] + idx = selidx[0] + select[idx] = 1 + selidx = torch.masked_select(selidx, torch.abs(ref - ntime) > self.min_gap) + nprob = torch.masked_select(nprob, torch.abs(ref - ntime) > self.min_gap) + ntime = torch.masked_select(ntime, torch.abs(ref - ntime) > self.min_gap) + p_time = torch.masked_select(time_sel[order], select > 0.0) + p_prob = torch.masked_select(score[order], select > 0.0) + p_type = torch.ones_like(p_time) * itr + y = torch.stack([p_type, p_time, p_prob], dim=1) + output.append(y) + y = torch.cat(output, dim=0) + return y diff --git a/makejit.eqt.py b/makejit.eqt.py index e36e337..b37e24c 100755 --- a/makejit.eqt.py +++ b/makejit.eqt.py @@ -1,73 +1,15 @@ -import torch -import torch.nn as nn -from models.EQT import EQTransformer -class Picker(EQTransformer): +import torch +from models.EQT import EQTransformer +from jit_picker_base import SlidingWindowPicker + + +class Picker(SlidingWindowPicker): def __init__(self): - super().__init__() - self.n_stride = 1 - def forward(self, x): - device = x.device - with torch.no_grad(): - #print("数据维度", x.shape) - T, C = x.shape - seqlen = 6144 - batchstride = 6144 - 256 - batchlen = torch.ceil(torch.tensor(T / batchstride).to(device)) - idx = torch.arange(0, seqlen, 1, device=device).unsqueeze(0) + torch.arange(0, batchlen, 1, device=device).unsqueeze(1) * batchstride - idx = idx.clamp(min=0, max=T-1).long() - x = x.to(device) - wave = x[idx, :] - wave = wave.permute(0, 2, 1) - wave -= torch.mean(wave, dim=2, keepdim=True) - max, maxidx = torch.max(torch.abs(wave), dim=2, keepdim=True) - wave /= (max + 1e-6) - x = self.encoder1(wave) - x = self.encoder2(x) - x = self.encoder3(x) - e = self.trans1(x) - y = self.decoder1(e) - oc = y.softmax(dim=1) - B, C, T = oc.shape - tgrid = torch.arange(0, T, 1, device=device).unsqueeze(0) * self.n_stride + torch.arange(0, batchlen, 1, device=device).unsqueeze(1) * batchstride - oc = oc.permute(0, 2, 1).reshape(-1, C) - ot = tgrid.squeeze() - ot = ot.reshape(-1) - output = [] - #print("NN处理完成", oc.shape, ot.shape) - # 接近非极大值抑制（NMS） - # .......P........S...... - for itr in range(2): - pc = oc[:, itr+1] - time_sel = torch.masked_select(ot, pc>0.3) - score = torch.masked_select(pc, pc>0.3) - _, order = score.sort(0, descending=True) # 降序排列 - ntime = time_sel[order] - nprob = score[order] - #print(batchstride, ntime, nprob) - select = -torch.ones_like(order) - selidx = torch.arange(0, order.numel(), 1, dtype=torch.long, device=device) - count = 0 - while True: - if nprob.numel()<1: - break - ref = ntime[0] - idx = selidx[0] - select[idx] = 1 - count += 1 - selidx = torch.masked_select(selidx, torch.abs(ref-ntime)>1000) - nprob = torch.masked_select(nprob, torch.abs(ref-ntime)>1000) - ntime = torch.masked_select(ntime, torch.abs(ref-ntime)>1000) - p_time = torch.masked_select(time_sel[order], select>0.0) - p_prob = torch.masked_select(score[order], select>0.0) - p_type = torch.ones_like(p_time) * itr - y = torch.stack([p_type, p_time, p_prob], dim=1) - output.append(y) - y = torch.cat(output, dim=0) - return y + super().__init__(EQTransformer, ckpt_path="ckpt/china.eqt.pt") + -model = Picker() -model.load_state_dict(torch.load("ckpt/china.eqt.pt", map_location="cpu")) +model = Picker() model.eval() torch.jit.save(torch.jit.script(model), "pickers/eqt.jit") x = torch.randn([300000, 3]) -y = model(x) \ No newline at end of file +y = model(x) diff --git a/makejit.eqtransformer.py b/makejit.eqtransformer.py index cb0e819..5421490 100644 --- a/makejit.eqtransformer.py +++ b/makejit.eqtransformer.py @@ -1,727 +1,20 @@ import torch -import seisbench -from seisbench.models import EQTransformer - -import warnings -from typing import Any - -import numpy as np -import torch -import torch.nn as nn -import torch.nn.functional as F - - - -# For implementation, potentially follow: https://medium.com/huggingface/from-tensorflow-to-pytorch-265f40ef2a28 -class EQTransformerForPicker(nn.Module): - """ - The EQTransformer from Mousavi et al. (2020) - - Implementation adapted from the Github repository https://github.com/smousavi05/EQTransformer - Assumes padding="same" and activation="relu" as in the pretrained EQTransformer models - - By instantiating the model with `from_pretrained("original")` a binary compatible version of the original - EQTransformer with the original weights from Mousavi et al. (2020) can be loaded. - - .. document_args:: seisbench.models EQTransformer - - :param in_channels: Number of input channels, by default 3. - :param in_samples: Number of input samples per channel, by default 6000. - The model expects input shape (in_channels, in_samples) - :param classes: Number of output classes, by default 2. The detection channel is not counted. - :param phases: Phase hints for the classes, by default "PS". Can be None. - :param res_cnn_blocks: Number of residual convolutional blocks - :param lstm_blocks: Number of LSTM blocks - :param drop_rate: Dropout rate - :param original_compatible: If True, uses a few custom layers for binary compatibility with original model - from Mousavi et al. (2020). - This option defaults to False. - It is usually recommended to stick to the default value, as the custom layers show - slightly worse performance than the PyTorch builtins. - The exception is when loading the original weights using :py:func:`from_pretrained`. - :param norm: Data normalization strategy, either "peak" or "std". - :param kwargs: Keyword arguments passed to the constructor of :py:class:`WaveformModel`. - """ - - - - def __init__( - self, - in_channels=3, - in_samples=6000, - classes=2, - phases="PS", - lstm_blocks=3, - drop_rate=0.1, - original_compatible=False, - sampling_rate=100, - norm="std", - **kwargs, - ): - citation = ( - "Mousavi, S.M., Ellsworth, W.L., Zhu, W., Chuang, L, Y., and Beroza, G, C. " - "Earthquake transformer—an attentive deep-learning model for simultaneous earthquake " - "detection and phase picking. Nat Commun 11, 3952 (2020). " - "https://doi.org/10.1038/s41467-020-17591-w" - ) - - # PickBlue options - for option in ("norm_amp_per_comp", "norm_detrend"): - if option in kwargs: - setattr(self, option, kwargs[option]) - del kwargs[option] - else: - setattr(self, option, False) - - # Blinding defines how many samples at beginning and end of the prediction should be ignored - # This is usually required to mitigate prediction problems from training properties, e.g., - # if all picks in the training fall between seconds 5 and 55. - super().__init__() - - self.in_channels = in_channels - self.classes = classes - self.lstm_blocks = lstm_blocks - self.drop_rate = drop_rate - self.norm = norm - self.in_samples = in_samples - - # Add options for conservative and the true original - see https://github.com/seisbench/seisbench/issues/96#issuecomment-1155158224 - if original_compatible: - warnings.warn( - "Using the non-conservative 'original' model, set `original_compatible='conservative' to use the more conservative model" - ) - original_compatible = "non-conservative" - - if original_compatible: - eps = 1e-7 # See Issue #96 - original models use tensorflow default epsilon of 1e-7 - else: - eps = 1e-5 - self.original_compatible = original_compatible - - if original_compatible and in_samples != 6000: - raise ValueError("original_compatible=True requires in_samples=6000.") - - self._phases = phases - if phases is not None and len(phases) != classes: - raise ValueError( - f"Number of classes ({classes}) does not match number of phases ({len(phases)})." - ) - - # Parameters from EQTransformer repository - self.filters = [ - 8, - 16, - 16, - 32, - 32, - 64, - 64, - ] # Number of filters for the convolutions - self.kernel_sizes = [11, 9, 7, 7, 5, 5, 3] # Kernel sizes for the convolutions - self.res_cnn_kernels = [3, 3, 3, 3, 2, 3, 2] - - # TODO: Add regularizers when training model - # kernel_regularizer=keras.regularizers.l2(1e-6), - # bias_regularizer=keras.regularizers.l1(1e-4), - - # Encoder stack - self.encoder = Encoder( - input_channels=self.in_channels, - filters=self.filters, - kernel_sizes=self.kernel_sizes, - in_samples=self.in_samples, - ) - - # Res CNN Stack - self.res_cnn_stack = ResCNNStack( - kernel_sizes=self.res_cnn_kernels, - filters=self.filters[-1], - drop_rate=self.drop_rate, - ) - - # BiLSTM stack - self.bi_lstm_stack = BiLSTMStack( - blocks=self.lstm_blocks, - input_size=self.filters[-1], - drop_rate=self.drop_rate, - original_compatible=original_compatible, - ) - - # Global attention - two transformers - self.transformer_d0 = Transformer( - input_size=16, drop_rate=self.drop_rate, eps=eps - ) - self.transformer_d = Transformer( - input_size=16, drop_rate=self.drop_rate, eps=eps - ) - - # Detection decoder and final Conv - self.decoder_d = Decoder( - input_channels=16, - filters=self.filters[::-1], - kernel_sizes=self.kernel_sizes[::-1], - out_samples=in_samples, - original_compatible=original_compatible, - ) - self.conv_d = nn.Conv1d( - in_channels=self.filters[0], out_channels=1, kernel_size=11, padding=5 - ) - - # Picking branches - self.pick_lstms = [] - self.pick_attentions = [] - self.pick_decoders = [] - self.pick_convs = [] - self.dropout = nn.Dropout(drop_rate) - - for _ in range(self.classes): - if original_compatible == "conservative": - # The non-conservative model uses a sigmoid activiation as handled by the base nn.LSTM - lstm = CustomLSTM(ActivationLSTMCell, 16, 16, bidirectional=False) - else: - lstm = nn.LSTM(16, 16, bidirectional=False) - self.pick_lstms.append(lstm) - - attention = SeqSelfAttention(input_size=16, attention_width=3, eps=eps) - self.pick_attentions.append(attention) - - decoder = Decoder( - input_channels=16, - filters=self.filters[::-1], - kernel_sizes=self.kernel_sizes[::-1], - out_samples=in_samples, - original_compatible=original_compatible, - ) - self.pick_decoders.append(decoder) - - conv = nn.Conv1d( - in_channels=self.filters[0], out_channels=1, kernel_size=11, padding=5 - ) - self.pick_convs.append(conv) - - self.pick_lstms = nn.ModuleList(self.pick_lstms) - self.pick_attentions = nn.ModuleList(self.pick_attentions) - self.pick_decoders = nn.ModuleList(self.pick_decoders) - self.pick_convs = nn.ModuleList(self.pick_convs) - - def forward(self, x): - assert x.ndim == 3 - assert x.shape[1:] == (self.in_channels, self.in_samples) - - # Shared encoder part - x = self.encoder(x) - x = self.res_cnn_stack(x) - x = self.bi_lstm_stack(x) - x, _ = self.transformer_d0(x) - x, _ = self.transformer_d(x) - - # Detection part - detection = self.decoder_d(x) - - detection = torch.sigmoid(self.conv_d(detection)) - detection = torch.squeeze(detection, dim=1) # Remove channel dimension - - outputs = [detection] - - # Pick parts - for lstm, attention, decoder, conv in zip( - self.pick_lstms, self.pick_attentions, self.pick_decoders, self.pick_convs - ): - px = x.permute( - 2, 0, 1 - ) # From batch, channels, sequence to sequence, batch, channels - px = lstm(px)[0] - px = self.dropout(px) - px = px.permute( - 1, 2, 0 - ) # From sequence, batch, channels to batch, channels, sequence - px, _ = attention(px) - px = decoder(px) - - pred = torch.sigmoid(conv(px)) - pred = torch.squeeze(pred, dim=1) # Remove channel dimension - - outputs.append(pred) - - return tuple(outputs) - - - - -class Encoder(nn.Module): - """ - Encoder stack - """ - - def __init__(self, input_channels, filters, kernel_sizes, in_samples): - super().__init__() - - convs = [] - pools = [] - self.paddings = [] - for in_channels, out_channels, kernel_size in zip( - [input_channels] + filters[:-1], filters, kernel_sizes - ): - convs.append( - nn.Conv1d( - in_channels, out_channels, kernel_size, padding=kernel_size // 2 - ) - ) - - # To be consistent with the behaviour in tensorflow, - # padding needs to be added for odd numbers of input_samples - padding = in_samples % 2 - - # Padding for MaxPool1d needs to be handled manually to conform with tf padding - self.paddings.append(padding) - pools.append(nn.MaxPool1d(2, padding=0)) - in_samples = (in_samples + padding) // 2 - - self.convs = nn.ModuleList(convs) - self.pools = nn.ModuleList(pools) - #self.paddings = nn.ModuleList(self.paddings) - - def forward(self, x): - # Use zip for convs and pools, access paddings by index - for i, (conv, pool) in enumerate(zip(self.convs, self.pools)): - padding = self.paddings[i] # Access padding by index - - x = torch.relu(conv(x)) # Apply convolution followed by ReLU - if padding != 0: - # Only pad right, use -1e10 as negative infinity - x = F.pad(x, (0, padding), "constant", -1000000.0) # Apply padding to the right - x = pool(x) # Apply pooling - - return x - - -class Decoder(nn.Module): - def __init__( - self, - input_channels, - filters, - kernel_sizes, - out_samples, - original_compatible=False, - ): - super().__init__() - - self.upsample = nn.Upsample(scale_factor=2, mode="nearest") - self.original_compatible = original_compatible - - # We need to trim off the final sample sometimes to get to the right number of output samples - self.crops = [] - current_samples = out_samples - for i, _ in enumerate(filters): - padding = current_samples % 2 - current_samples = (current_samples + padding) // 2 - if padding == 1: - self.crops.append(len(filters) - 1 - i) - - convs = [] - for in_channels, out_channels, kernel_size in zip( - [input_channels] + filters[:-1], filters, kernel_sizes - ): - convs.append( - nn.Conv1d( - in_channels, out_channels, kernel_size, padding=kernel_size // 2 - ) - ) - - self.convs = nn.ModuleList(convs) - - def forward(self, x): - for i, conv in enumerate(self.convs): - x = self.upsample(x) - - if self.original_compatible: - if i == 3: - x = x[:, :, 1:-1] - else: - if i in self.crops: - x = x[:, :, :-1] - - x = F.relu(conv(x)) - - return x - - -class ResCNNStack(nn.Module): - def __init__(self, kernel_sizes, filters, drop_rate): - super().__init__() - - members = [] - for ker in kernel_sizes: - members.append(ResCNNBlock(filters, ker, drop_rate)) - - self.members = nn.ModuleList(members) - - def forward(self, x): - for member in self.members: - x = member(x) - - return x - - -class ResCNNBlock(nn.Module): - def __init__(self, filters, ker, drop_rate): - super().__init__() - - self.manual_padding = False - if ker == 3: - padding = 1 - else: - # ker == 2 - # Manual padding emulate the padding in tensorflow - self.manual_padding = True - padding = 0 - - self.dropout = SpatialDropout1d(drop_rate) - - self.norm1 = nn.BatchNorm1d(filters, eps=1e-3) - self.conv1 = nn.Conv1d(filters, filters, ker, padding=padding) - - self.norm2 = nn.BatchNorm1d(filters, eps=1e-3) - self.conv2 = nn.Conv1d(filters, filters, ker, padding=padding) - - def forward(self, x): - y = self.norm1(x) - y = F.relu(y) - y = self.dropout(y) - if self.manual_padding: - y = F.pad(y, (0, 1), "constant", 0.0) - y = self.conv1(y) - - y = self.norm2(y) - y = F.relu(y) - y = self.dropout(y) - if self.manual_padding: - y = F.pad(y, (0, 1), "constant", 0.0) - y = self.conv2(y) - - return x + y - - -class BiLSTMStack(nn.Module): - def __init__( - self, blocks, input_size, drop_rate, hidden_size=16, original_compatible=False - ): - super().__init__() - - # First LSTM has a different input size as the subsequent ones - self.members = nn.ModuleList( - [ - BiLSTMBlock( - input_size, - hidden_size, - drop_rate, - original_compatible=original_compatible, - ) - ] - + [ - BiLSTMBlock( - hidden_size, - hidden_size, - drop_rate, - original_compatible=original_compatible, - ) - for _ in range(blocks - 1) - ] - ) - - def forward(self, x): - for member in self.members: - x = member(x) - return x - - -class BiLSTMBlock(nn.Module): - def __init__(self, input_size, hidden_size, drop_rate, original_compatible=False): - super().__init__() - - if original_compatible == "conservative": - # The non-conservative model uses a sigmoid activiation as handled by the base nn.LSTM - self.lstm = CustomLSTM(ActivationLSTMCell, input_size, hidden_size) - elif original_compatible == "non-conservative": - self.lstm = CustomLSTM( - ActivationLSTMCell, - input_size, - hidden_size, - gate_activation=torch.sigmoid, - ) - else: - self.lstm = nn.LSTM(input_size, hidden_size, bidirectional=True) - self.dropout = nn.Dropout(drop_rate) - self.conv = nn.Conv1d(2 * hidden_size, hidden_size, 1) - self.norm = nn.BatchNorm1d(hidden_size, eps=1e-3) - - def forward(self, x): - x = x.permute( - 2, 0, 1 - ) # From batch, channels, sequence to sequence, batch, channels - x = self.lstm(x)[0] - x = self.dropout(x) - x = x.permute( - 1, 2, 0 - ) # From sequence, batch, channels to batch, channels, sequence - x = self.conv(x) - x = self.norm(x) - return x - - -class Transformer(nn.Module): - def __init__(self, input_size, drop_rate, attention_width=None, eps=1e-5): - super().__init__() - - self.attention = SeqSelfAttention( - input_size, attention_width=attention_width, eps=eps - ) - self.norm1 = LayerNormalization(input_size) - self.ff = FeedForward(input_size, drop_rate) - self.norm2 = LayerNormalization(input_size) - - def forward(self, x): - y, weight = self.attention(x) - y = x + y - y = self.norm1(y) - y2 = self.ff(y) - y2 = y + y2 - y2 = self.norm2(y2) - - return y2, weight - - -class SeqSelfAttention(nn.Module): - """ - Additive self attention - """ - - def __init__(self, input_size, units=32, attention_width=None, eps=1e-5): - super().__init__() - self.attention_width = attention_width - - self.Wx = nn.Parameter(uniform(-0.02, 0.02, input_size, units)) - self.Wt = nn.Parameter(uniform(-0.02, 0.02, input_size, units)) - self.bh = nn.Parameter(torch.zeros(units)) - - self.Wa = nn.Parameter(uniform(-0.02, 0.02, units, 1)) - self.ba = nn.Parameter(torch.zeros(1)) - - self.eps = eps - - def forward(self, x): - # x.shape == (batch, channels, time) - - x = x.permute(0, 2, 1) # to (batch, time, channels) - - q = torch.unsqueeze( - torch.matmul(x, self.Wt), 2 - ) # Shape (batch, time, 1, channels) - k = torch.unsqueeze( - torch.matmul(x, self.Wx), 1 - ) # Shape (batch, 1, time, channels) - - h = torch.tanh(q + k + self.bh) - - # Emissions - e = torch.squeeze( - torch.matmul(h, self.Wa) + self.ba, -1 - ) # Shape (batch, time, time) - - # This is essentially softmax with an additional attention component. - e = ( - e - torch.max(e, dim=-1, keepdim=True).values - ) # In versions <= 0.2.1 e was incorrectly normalized by max(x) - e = torch.exp(e) - if self.attention_width is not None: - lower = ( - torch.arange(0, e.shape[1], device=e.device) - self.attention_width // 2 - ) - upper = lower + self.attention_width - indices = torch.unsqueeze(torch.arange(0, e.shape[1], device=e.device), 1) - mask = torch.logical_and(lower <= indices, indices < upper) - e = torch.where(mask, e, torch.zeros_like(e)) - - a = e / (torch.sum(e, dim=-1, keepdim=True) + self.eps) - - v = torch.matmul(a, x) - - v = v.permute(0, 2, 1) # to (batch, channels, time) - - return v, a - - -def uniform(a, b, *args): - return a + (b - a) * torch.rand(*args) - - -class LayerNormalization(nn.Module): - def __init__(self, filters, eps=1e-14): - super().__init__() - - gamma = torch.ones(filters, 1) - self.gamma = nn.Parameter(gamma) - beta = torch.zeros(filters, 1) - self.beta = nn.Parameter(beta) - self.eps = eps - - def forward(self, x): - mean = torch.mean(x, 1, keepdim=True) - var = torch.mean((x - mean) ** 2, 1, keepdim=True) + self.eps - std = torch.sqrt(var) - outputs = (x - mean) / std - - outputs = outputs * self.gamma - outputs = outputs + self.beta - - return outputs - - -class FeedForward(nn.Module): - def __init__(self, io_size, drop_rate, hidden_size=128): - super().__init__() - - self.lin1 = nn.Linear(io_size, hidden_size) - self.lin2 = nn.Linear(hidden_size, io_size) - self.dropout = nn.Dropout(drop_rate) - - def forward(self, x): - x = x.permute(0, 2, 1) # To (batch, time, channel) - x = F.relu(self.lin1(x)) - x = self.dropout(x) - x = self.lin2(x) - x = x.permute(0, 2, 1) # To (batch, channel, time) - - return x - - -class SpatialDropout1d(nn.Module): - def __init__(self, drop_rate): - super().__init__() - - self.drop_rate = drop_rate - self.dropout = nn.Dropout2d(drop_rate) - - def forward(self, x): - x = x.unsqueeze(dim=-1) # Add fake dimension - x = self.dropout(x) - x = x.squeeze(dim=-1) # Remove fake dimension - return x - +from models.EQTransformer import EQTransformer +from jit_picker_base import SlidingWindowPicker +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +pretrained = EQTransformer.from_pretrained("stead").to(device) +pretrained_state = {k: v.cpu() for k, v in pretrained.state_dict().items()} -class Picker(EQTransformerForPicker): +class Picker(SlidingWindowPicker): def __init__(self): - super().__init__() - - def forward(self, x): - device = x.device - with torch.no_grad(): - #print("数据维度", x.shape) - T, C = x.shape - seqlen = 6000 - batchstride = seqlen - 3000 - batchlen = torch.ceil(torch.tensor(T / batchstride).to(device)) - idx = torch.arange(0, seqlen, 1, device=device).unsqueeze(0) + torch.arange(0, batchlen, 1, device=device).unsqueeze(1) * batchstride - idx = idx.clamp(min=0, max=T-1).long() - x = x.to(device) - wave = x[idx, :] - wave = wave.permute(0, 2, 1) - wave -= torch.mean(wave, dim=2, keepdim=True) - #max, maxidx = torch.max(torch.abs(wave), dim=2, keepdim=True) - #max, maxidx = torch.max(max, dim=1, keepdim=True) - max = torch.std(wave, dim=2, keepdim=True) - wave /= (max + 1e-6) - x = wave - x = self.encoder(x) - x = self.res_cnn_stack(x) - x = self.bi_lstm_stack(x) - x, _ = self.transformer_d0(x) - x, _ = self.transformer_d(x) - - # Detection part - detection = self.decoder_d(x) - - detection = torch.sigmoid(self.conv_d(detection)) - detection = torch.squeeze(detection, dim=1) # Remove channel dimension - - outputs = [detection] - - # Pick parts - for lstm, attention, decoder, conv in zip( - self.pick_lstms, self.pick_attentions, self.pick_decoders, self.pick_convs - ): - px = x.permute( - 2, 0, 1 - ) # From batch, channels, sequence to sequence, batch, channels - px = lstm(px)[0] - px = self.dropout(px) - px = px.permute( - 1, 2, 0 - ) # From sequence, batch, channels to batch, channels, sequence - px, _ = attention(px) - px = decoder(px) - - pred = torch.sigmoid(conv(px)) - pred = torch.squeeze(pred, dim=1) # Remove channel dimension - - outputs.append(pred) - oc = torch.stack(outputs, dim=1) - B, C, T = oc.shape - tgrid = torch.arange(0, T, 1, device=device).unsqueeze(0) * 1 + torch.arange(0, batchlen, 1, device=device).unsqueeze(1) * batchstride - oc = oc.permute(0, 2, 1).reshape(-1, C) - #oc = oc[:, [2, 0, 1]] - ot = tgrid.squeeze() - ot = ot.reshape(-1) - output = [] - #print("NN处理完成", oc.shape, ot.shape) - # 接近非极大值抑制（NMS） - # .......P........S...... - #oc = oc.cpu() - #ot = ot.cpu() - for itr in range(2): - pc = oc[:, itr+1] - time_sel = torch.masked_select(ot, pc>0.1) - score = torch.masked_select(pc, pc>0.1) - _, order = score.sort(0, descending=True) # 降序排列 - ntime = time_sel[order] - nprob = score[order] - #print(batchstride, ntime, nprob) - select = -torch.ones_like(order) - selidx = torch.arange(0, order.numel(), 1, dtype=torch.long, device=device) - count = 0 - while True: - if nprob.numel()<1: - break - ref = ntime[0] - idx = selidx[0] - select[idx] = 1 - count += 1 - selidx = torch.masked_select(selidx, torch.abs(ref-ntime)>1000) - nprob = torch.masked_select(nprob, torch.abs(ref-ntime)>1000) - ntime = torch.masked_select(ntime, torch.abs(ref-ntime)>1000) - p_time = torch.masked_select(time_sel[order], select>0.0) - p_prob = torch.masked_select(score[order], select>0.0) - p_type = torch.ones_like(p_time) * itr - y = torch.stack([p_type, p_time, p_prob], dim=1) - output.append(y) - y = torch.cat(output, dim=0) - return y - - - - -# Define the device (use 'cuda' if you have a GPU, else 'cpu') -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + super().__init__(EQTransformer, state_dict=pretrained_state) -# Load the pre-trained EqTransformer model from SeisBench (you can replace with your own model path) -model2 = EQTransformer.from_pretrained('stead').to(device) -model = Picker() -model.load_state_dict(model2.state_dict()) +model = Picker() model.eval() torch.jit.save(torch.jit.script(model), "pickers/eqtransformer.stead.jit") x = torch.randn([300000, 3]) -y = model(x) \ No newline at end of file +y = model(x) diff --git a/makejit.pnsn.diff.py b/makejit.pnsn.diff.py index 216baf3..3bf8996 100755 --- a/makejit.pnsn.diff.py +++ b/makejit.pnsn.diff.py @@ -1,91 +1,15 @@ -import torch -import torch.nn as nn -from models.BRNNPNSN import BRNN -class Picker(BRNN): - def __init__(self): - super().__init__() - self.n_stride = 1 - def forward(self, x): - device = x.device - with torch.no_grad(): - #print("数据维度", x.shape) - T, C = x.shape - seqlen = 10240 - batchstride = 10240 - 512 - batchlen = torch.ceil(torch.tensor(T / batchstride).to(device)) - idx = torch.arange(0, seqlen, 1, device=device).unsqueeze(0) + torch.arange(0, batchlen, 1, device=device).unsqueeze(1) * batchstride - idx = idx.clamp(min=0, max=T-2).long() - x = x.to(device) - origx = x - - wave = x[idx, :] - wave = wave.permute(0, 2, 1) - wave -= torch.mean(wave, dim=2, keepdim=True) - max = torch.std(wave, dim=2, keepdim=True) - #max, maxidx = torch.max(torch.abs(wave), dim=2, keepdim=True) - wave /= (max + 1e-6) - #print(wave.shape) - x = self.encoder(wave) - e = self.rnns(x) # 波形特征 - y = self.decoder(e) # 输出概率 - oc1 = y.softmax(dim=1) +import torch +from models.BRNN import BRNN +from jit_picker_base import SlidingWindowPicker - wave = (origx[1:]-origx[:-1])[idx, :] - wave = wave.permute(0, 2, 1) - wave -= torch.mean(wave, dim=2, keepdim=True) - max = torch.std(wave, dim=2, keepdim=True) - #max, maxidx = torch.max(torch.abs(wave), dim=2, keepdim=True) - wave /= (max + 1e-6) - #print(wave.shape) - x = self.encoder(wave) - e = self.rnns(x) # 波形特征 - y = self.decoder(e) # 输出概率 - oc2 = y.softmax(dim=1) - oc = torch.cat([oc1, oc2], dim=0) +class Picker(SlidingWindowPicker): + def __init__(self): + super().__init__(BRNN, ckpt_path="ckpt/china.rnn.pnsn.pt") - B, C, T = oc.shape - tgrid = torch.arange(0, T, 1, device=device).unsqueeze(0) * self.n_stride + torch.arange(0, batchlen, 1, device=device).unsqueeze(1) * batchstride - oc = oc.permute(0, 2, 1).reshape(-1, C) - ot = tgrid.squeeze() - ot = ot.reshape(-1) - ot = torch.cat([ot, ot], dim=0) - output = [] - #print("NN处理完成", oc.shape, ot.shape) - # 接近非极大值抑制（NMS） - # .......P........S...... - for itr in range(4): - pc = oc[:, itr+1] - time_sel = torch.masked_select(ot, pc>0.1) - score = torch.masked_select(pc, pc>0.1) - _, order = score.sort(0, descending=True) # 降序排列 - ntime = time_sel[order] - nprob = score[order] - #print(batchstride, ntime, nprob) - select = -torch.ones_like(order) - selidx = torch.arange(0, order.numel(), 1, dtype=torch.long, device=device) - count = 0 - while True: - if nprob.numel()<1: - break - ref = ntime[0] - idx = selidx[0] - select[idx] = 1 - count += 1 - selidx = torch.masked_select(selidx, torch.abs(ref-ntime)>300) - nprob = torch.masked_select(nprob, torch.abs(ref-ntime)>300) - ntime = torch.masked_select(ntime, torch.abs(ref-ntime)>300) - p_time = torch.masked_select(time_sel[order], select>0.0) - p_prob = torch.masked_select(score[order], select>0.0) - p_type = torch.ones_like(p_time) * (itr) - y = torch.stack([p_type, p_time, p_prob], dim=1) - output.append(y) - y = torch.cat(output, dim=0) - return y -model = Picker() -model.load_state_dict(torch.load("ckpt/china.rnn.pnsn.pt", map_location="cpu")) +model = Picker() model.eval() torch.jit.save(torch.jit.script(model), "pickers/rnn.origdiff.01.jit") x = torch.randn([300000, 3]) -y = model(x) \ No newline at end of file +y = model(x) diff --git a/makejit.pnsn.py b/makejit.pnsn.py index be16638..fcaacb4 100755 --- a/makejit.pnsn.py +++ b/makejit.pnsn.py @@ -1,73 +1,15 @@ -import torch -import torch.nn as nn -from models.BRNNPNSN import BRNN -class Picker(BRNN): +import torch +from models.BRNN import BRNN +from jit_picker_base import SlidingWindowPicker + + +class Picker(SlidingWindowPicker): def __init__(self): - super().__init__() - self.n_stride = 1 - def forward(self, x): - device = x.device - with torch.no_grad(): - #print("数据维度", x.shape) - T, C = x.shape - seqlen = 10240 - batchstride = 10240 - 512 - batchlen = torch.ceil(torch.tensor(T / batchstride).to(device)) - idx = torch.arange(0, seqlen, 1, device=device).unsqueeze(0) + torch.arange(0, batchlen, 1, device=device).unsqueeze(1) * batchstride - idx = idx.clamp(min=0, max=T-1).long() - x = x.to(device) - wave = x[idx, :] - wave = wave.permute(0, 2, 1) - wave -= torch.mean(wave, dim=2, keepdim=True) - max = torch.std(wave, dim=2, keepdim=True) - #max, maxidx = torch.max(torch.abs(wave), dim=2, keepdim=True) - wave /= (max + 1e-6) - #print(wave.shape) - x = self.encoder(wave) - e = self.rnns(x) # 波形特征 - y = self.decoder(e) # 输出概率 - oc = y.softmax(dim=1) - B, C, T = oc.shape - tgrid = torch.arange(0, T, 1, device=device).unsqueeze(0) * self.n_stride + torch.arange(0, batchlen, 1, device=device).unsqueeze(1) * batchstride - oc = oc.permute(0, 2, 1).reshape(-1, C) - ot = tgrid.squeeze() - ot = ot.reshape(-1) - output = [] - #print("NN处理完成", oc.shape, ot.shape) - # 接近非极大值抑制（NMS） - # .......P........S...... - for itr in range(4): - pc = oc[:, itr+1] - time_sel = torch.masked_select(ot, pc>0.1) - score = torch.masked_select(pc, pc>0.1) - _, order = score.sort(0, descending=True) # 降序排列 - ntime = time_sel[order] - nprob = score[order] - #print(batchstride, ntime, nprob) - select = -torch.ones_like(order) - selidx = torch.arange(0, order.numel(), 1, dtype=torch.long, device=device) - count = 0 - while True: - if nprob.numel()<1: - break - ref = ntime[0] - idx = selidx[0] - select[idx] = 1 - count += 1 - selidx = torch.masked_select(selidx, torch.abs(ref-ntime)>300) - nprob = torch.masked_select(nprob, torch.abs(ref-ntime)>300) - ntime = torch.masked_select(ntime, torch.abs(ref-ntime)>300) - p_time = torch.masked_select(time_sel[order], select>0.0) - p_prob = torch.masked_select(score[order], select>0.0) - p_type = torch.ones_like(p_time) * itr - y = torch.stack([p_type, p_time, p_prob], dim=1) - output.append(y) - y = torch.cat(output, dim=0) - return y + super().__init__(BRNN, ckpt_path="ckpt/china.rnn.pnsn.pt") + -model = Picker() -model.load_state_dict(torch.load("ckpt/china.rnn.pnsn.pt", map_location="cpu")) +model = Picker() model.eval() -torch.jit.save(torch.jit.script(model), "pickers/rnn.pnsn.01.jit") +torch.jit.save(torch.jit.script(model), "pickers/pnsn.jit") x = torch.randn([300000, 3]) -y = model(x) \ No newline at end of file +y = model(x) diff --git a/makejit.rnn.py b/makejit.rnn.py index caab16a..a3ccfff 100755 --- a/makejit.rnn.py +++ b/makejit.rnn.py @@ -1,71 +1,15 @@ -import torch -import torch.nn as nn -from models.BRNN import BRNN -class Picker(BRNN): +import torch +from models.BRNN import BRNN +from jit_picker_base import SlidingWindowPicker + + +class Picker(SlidingWindowPicker): def __init__(self): - super().__init__() - self.n_stride = 1 - def forward(self, x): - device = x.device - with torch.no_grad(): - #print("数据维度", x.shape) - T, C = x.shape - seqlen = 6144 - batchstride = 6144 - 256 - batchlen = torch.ceil(torch.tensor(T / batchstride).to(device)) - idx = torch.arange(0, seqlen, 1, device=device).unsqueeze(0) + torch.arange(0, batchlen, 1, device=device).unsqueeze(1) * batchstride - idx = idx.clamp(min=0, max=T-1).long() - x = x.to(device) - wave = x[idx, :] - wave = wave.permute(0, 2, 1) - wave -= torch.mean(wave, dim=2, keepdim=True) - max, maxidx = torch.max(torch.abs(wave), dim=2, keepdim=True) - wave /= (max + 1e-6) - x = self.encoder(wave) - e = self.rnns(x) # 波形特征 - y = self.decoder(e) # 输出概率 - oc = y.softmax(dim=1) - B, C, T = oc.shape - tgrid = torch.arange(0, T, 1, device=device).unsqueeze(0) * self.n_stride + torch.arange(0, batchlen, 1, device=device).unsqueeze(1) * batchstride - oc = oc.permute(0, 2, 1).reshape(-1, C) - ot = tgrid.squeeze() - ot = ot.reshape(-1) - output = [] - #print("NN处理完成", oc.shape, ot.shape) - # 接近非极大值抑制（NMS） - # .......P........S...... - for itr in range(2): - pc = oc[:, itr+1] - time_sel = torch.masked_select(ot, pc>0.3) - score = torch.masked_select(pc, pc>0.3) - _, order = score.sort(0, descending=True) # 降序排列 - ntime = time_sel[order] - nprob = score[order] - #print(batchstride, ntime, nprob) - select = -torch.ones_like(order) - selidx = torch.arange(0, order.numel(), 1, dtype=torch.long, device=device) - count = 0 - while True: - if nprob.numel()<1: - break - ref = ntime[0] - idx = selidx[0] - select[idx] = 1 - count += 1 - selidx = torch.masked_select(selidx, torch.abs(ref-ntime)>1000) - nprob = torch.masked_select(nprob, torch.abs(ref-ntime)>1000) - ntime = torch.masked_select(ntime, torch.abs(ref-ntime)>1000) - p_time = torch.masked_select(time_sel[order], select>0.0) - p_prob = torch.masked_select(score[order], select>0.0) - p_type = torch.ones_like(p_time) * itr - y = torch.stack([p_type, p_time, p_prob], dim=1) - output.append(y) - y = torch.cat(output, dim=0) - return y + super().__init__(BRNN, ckpt_path="ckpt/china.rnn.pt") + -model = Picker() -model.load_state_dict(torch.load("ckpt/china.rnn.pt", map_location="cpu")) +model = Picker() model.eval() torch.jit.save(torch.jit.script(model), "pickers/rnn.jit") x = torch.randn([300000, 3]) -y = model(x) \ No newline at end of file +y = model(x) diff --git a/makejit.unet.py b/makejit.unet.py index 57e24b6..80fb327 100755 --- a/makejit.unet.py +++ b/makejit.unet.py @@ -1,85 +1,15 @@ -import torch -import torch.nn as nn -from models.UNet import UNet -class Picker(UNet): +import torch +from models.UNet import UNet +from jit_picker_base import SlidingWindowPicker + + +class Picker(SlidingWindowPicker): def __init__(self): - super().__init__() - self.n_stride = 1 - def forward(self, x): - device = x.device - with torch.no_grad(): - #print("数据维度", x.shape) - T, C = x.shape - seqlen = 6144 - batchstride = 6144 - 256 - batchlen = torch.ceil(torch.tensor(T / batchstride).to(device)) - idx = torch.arange(0, seqlen, 1, device=device).unsqueeze(0) + torch.arange(0, batchlen, 1, device=device).unsqueeze(1) * batchstride - idx = idx.clamp(min=0, max=T-1).long() - x = x.to(device) - wave = x[idx, :] - wave = wave.permute(0, 2, 1) - wave -= torch.mean(wave, dim=2, keepdim=True) - max, maxidx = torch.max(torch.abs(wave), dim=2, keepdim=True) - wave /= (max + 1e-6) - x = wave.unsqueeze(3) - x = self.inputs(x) - x1 = self.layer0(x) - x2 = self.layer1(x1) - x3 = self.layer2(x2) - x4 = self.layer3(x3) - x5 = self.layer4(x4) - x6 = self.layer5(x5) - x6 = torch.cat([x4, x6], dim=1) # 加入skip connection - x7 = self.layer6(x6) - x7 = torch.cat([x3, x7], dim=1) # 加入skip connection - x8 = self.layer7(x7) - x8 = torch.cat([x2, x8], dim=1) # 加入skip connection - x9 = self.layer8(x8) - x9 = torch.cat([x1, x9], dim=1) # 加入skip connection - x10 = self.layer9(x9) - x10 = x10.softmax(dim=1) - oc = x10.squeeze(dim=3) - B, C, T = oc.shape - tgrid = torch.arange(0, T, 1, device=device).unsqueeze(0) * self.n_stride + torch.arange(0, batchlen, 1, device=device).unsqueeze(1) * batchstride - oc = oc.permute(0, 2, 1).reshape(-1, C) - ot = tgrid.squeeze() - ot = ot.reshape(-1) - output = [] - #print("NN处理完成", oc.shape, ot.shape) - # 接近非极大值抑制（NMS） - # .......P........S...... - for itr in range(2): - pc = oc[:, itr+1] - time_sel = torch.masked_select(ot, pc>0.3) - score = torch.masked_select(pc, pc>0.3) - _, order = score.sort(0, descending=True) # 降序排列 - ntime = time_sel[order] - nprob = score[order] - #print(batchstride, ntime, nprob) - select = -torch.ones_like(order) - selidx = torch.arange(0, order.numel(), 1, dtype=torch.long, device=device) - count = 0 - while True: - if nprob.numel()<1: - break - ref = ntime[0] - idx = selidx[0] - select[idx] = 1 - count += 1 - selidx = torch.masked_select(selidx, torch.abs(ref-ntime)>1000) - nprob = torch.masked_select(nprob, torch.abs(ref-ntime)>1000) - ntime = torch.masked_select(ntime, torch.abs(ref-ntime)>1000) - p_time = torch.masked_select(time_sel[order], select>0.0) - p_prob = torch.masked_select(score[order], select>0.0) - p_type = torch.ones_like(p_time) * itr - y = torch.stack([p_type, p_time, p_prob], dim=1) - output.append(y) - y = torch.cat(output, dim=0) - return y + super().__init__(UNet, ckpt_path="ckpt/china.unet.pt") + -model = Picker() -model.load_state_dict(torch.load("ckpt/china.unet.pt", map_location="cpu")) +model = Picker() model.eval() torch.jit.save(torch.jit.script(model), "pickers/unet.jit") x = torch.randn([300000, 3]) -y = model(x) \ No newline at end of file +y = model(x) diff --git a/makejit.unetpp.py b/makejit.unetpp.py index d37a0ea..2760ad9 100755 --- a/makejit.unetpp.py +++ b/makejit.unetpp.py @@ -1,89 +1,15 @@ -import torch -import torch.nn as nn -from models.UNetPlusPlus import UNetpp -class Picker(UNetpp): - def __init__(self): - super().__init__() - self.n_stride = 1 - def forward(self, x): - device = x.device - with torch.no_grad(): - #print("数据维度", x.shape) - T, C = x.shape - seqlen = 6144 - batchstride = 6144 - 256 - batchlen = torch.ceil(torch.tensor(T / batchstride).to(device)) - idx = torch.arange(0, seqlen, 1, device=device).unsqueeze(0) + torch.arange(0, batchlen, 1, device=device).unsqueeze(1) * batchstride - idx = idx.clamp(min=0, max=T-1).long() - x = x.to(device) - wave = x[idx, :] - wave = wave.permute(0, 2, 1) - wave -= torch.mean(wave, dim=2, keepdim=True) - max, maxidx = torch.max(torch.abs(wave), dim=2, keepdim=True) - wave /= (max + 1e-6) - #x = wave.unsqueeze(3) - x0_0 = self.conv0_0(wave) - x1_0 = self.conv1_0(self.pool(x0_0)) - x0_1 = self.conv0_1(torch.cat([x0_0, self.up(x1_0)], 1)) - - x2_0 = self.conv2_0(self.pool(x1_0)) - x1_1 = self.conv1_1(torch.cat([x1_0, self.up(x2_0)], 1)) - x0_2 = self.conv0_2(torch.cat([x0_0, x0_1, self.up(x1_1)], 1)) +import torch +from models.UNetPlusPlus import UNetpp +from jit_picker_base import SlidingWindowPicker - x3_0 = self.conv3_0(self.pool(x2_0)) - x2_1 = self.conv2_1(torch.cat([x2_0, self.up(x3_0)], 1)) - x1_2 = self.conv1_2(torch.cat([x1_0, x1_1, self.up(x2_1)], 1)) - x0_3 = self.conv0_3(torch.cat([x0_0, x0_1, x0_2, self.up(x1_2)], 1)) - x4_0 = self.conv4_0(self.pool(x3_0)) - x3_1 = self.conv3_1(torch.cat([x3_0, self.up(x4_0)], 1)) - x2_2 = self.conv2_2(torch.cat([x2_0, x2_1, self.up(x3_1)], 1)) - x1_3 = self.conv1_3(torch.cat([x1_0, x1_1, x1_2, self.up(x2_2)], 1)) - x0_4 = self.conv0_4(torch.cat([x0_0, x0_1, x0_2, x0_3, self.up(x1_3)], 1)) +class Picker(SlidingWindowPicker): + def __init__(self): + super().__init__(UNetpp, ckpt_path="ckpt/china.unetpp.pt") - output = self.final(x0_4) - oc = self.sigmoid(output) - B, C, T = oc.shape - tgrid = torch.arange(0, T, 1, device=device).unsqueeze(0) * self.n_stride + torch.arange(0, batchlen, 1, device=device).unsqueeze(1) * batchstride - oc = oc.permute(0, 2, 1).reshape(-1, C) - ot = tgrid.squeeze() - ot = ot.reshape(-1) - output = [] - #print("NN处理完成", oc.shape, ot.shape) - # 接近非极大值抑制（NMS） - # .......P........S...... - for itr in range(2): - pc = oc[:, itr+1] - time_sel = torch.masked_select(ot, pc>0.3) - score = torch.masked_select(pc, pc>0.3) - _, order = score.sort(0, descending=True) # 降序排列 - ntime = time_sel[order] - nprob = score[order] - #print(batchstride, ntime, nprob) - select = -torch.ones_like(order) - selidx = torch.arange(0, order.numel(), 1, dtype=torch.long, device=device) - count = 0 - while True: - if nprob.numel()<1: - break - ref = ntime[0] - idx = selidx[0] - select[idx] = 1 - count += 1 - selidx = torch.masked_select(selidx, torch.abs(ref-ntime)>1000) - nprob = torch.masked_select(nprob, torch.abs(ref-ntime)>1000) - ntime = torch.masked_select(ntime, torch.abs(ref-ntime)>1000) - p_time = torch.masked_select(time_sel[order], select>0.0) - p_prob = torch.masked_select(score[order], select>0.0) - p_type = torch.ones_like(p_time) * itr - y = torch.stack([p_type, p_time, p_prob], dim=1) - output.append(y) - y = torch.cat(output, dim=0) - return y -model = Picker() -model.load_state_dict(torch.load("ckpt/china.unetpp.pt", map_location="cpu")) +model = Picker() model.eval() torch.jit.save(torch.jit.script(model), "pickers/unetpp.jit") x = torch.randn([300000, 3]) -y = model(x) \ No newline at end of file +y = model(x) diff --git a/makeonnx.eqt.py b/makeonnx.eqt.py index 340f3ed..16a4faf 100755 --- a/makeonnx.eqt.py +++ b/makeonnx.eqt.py @@ -1,46 +1,25 @@ -import torch -import torch.nn as nn -from models.EQT import EQTransformer -class Picker(EQTransformer): +import torch +from models.EQT import EQTransformer +from onnx_picker_base import OnnxSlidingWindowPicker + + +class Picker(OnnxSlidingWindowPicker): def __init__(self): - super().__init__() - self.n_stride = 1 - def forward(self, x): - device = x.device - with torch.no_grad(): - #print("数据维度", x.shape) - T, C = x.shape - seqlen = 6144 - batchstride = 6144 - 256 - batchlen = torch.ceil(torch.tensor(T / batchstride).to(device)) - idx = torch.arange(0, seqlen, 1, device=device).unsqueeze(0) + torch.arange(0, batchlen, 1, device=device).unsqueeze(1) * batchstride - idx = idx.clamp(min=0, max=T-1).long() - x = x.to(device) - wave = x[idx, :] - wave = wave.permute(0, 2, 1) - wave -= torch.mean(wave, dim=2, keepdim=True) - max, maxidx = torch.max(torch.abs(wave), dim=2, keepdim=True) - wave /= (max + 1e-6) - x = self.encoder1(wave) - x = self.encoder2(x) - x = self.encoder3(x) - e = self.trans1(x) - y = self.decoder1(e) - oc = y.softmax(dim=1) - B, C, T = oc.shape - tgrid = torch.arange(0, T, 1, device=device).unsqueeze(0) * self.n_stride + torch.arange(0, batchlen, 1, device=device).unsqueeze(1) * batchstride - oc = oc.permute(0, 2, 1).reshape(-1, C) - ot = tgrid.squeeze() - ot = ot.reshape(-1) - return oc, ot -model = Picker() + super().__init__(EQTransformer, ckpt_path="ckpt/china.eqt.pt") + + +model = Picker() model.eval() -model.load_state_dict(torch.load("ckpt/china.eqt.pt", map_location="cpu")) input_names = ["wave"] output_names = ["prob", "time"] -#x = torch.randn([10, 3, 6144, 1]) x = torch.randn([500000, 3]) -torch.onnx.export(model, x, -"pickers/eqt.onnx", verbose=True, -dynamic_axes={"wave":{0:"batch"}, "prob":{0:"batch"}, "time":{0:"batch"}}, -input_names=input_names, output_names=output_names, opset_version=11) \ No newline at end of file +torch.onnx.export( + model, + x, + "pickers/eqt.onnx", + verbose=True, + dynamic_axes={"wave": {0: "batch"}, "prob": {0: "batch"}, "time": {0: "batch"}}, + input_names=input_names, + output_names=output_names, + opset_version=11, +) diff --git a/makeonnx.eqtransformer.py b/makeonnx.eqtransformer.py index 99f1259..367ef10 100644 --- a/makeonnx.eqtransformer.py +++ b/makeonnx.eqtransformer.py @@ -1,700 +1,30 @@ import torch -import seisbench -from seisbench.models import EQTransformer - -import warnings -from typing import Any - -import numpy as np -import torch -import torch.nn as nn -import torch.nn.functional as F - - - -# For implementation, potentially follow: https://medium.com/huggingface/from-tensorflow-to-pytorch-265f40ef2a28 -class EQTransformerForPicker(nn.Module): - """ - The EQTransformer from Mousavi et al. (2020) - - Implementation adapted from the Github repository https://github.com/smousavi05/EQTransformer - Assumes padding="same" and activation="relu" as in the pretrained EQTransformer models - - By instantiating the model with `from_pretrained("original")` a binary compatible version of the original - EQTransformer with the original weights from Mousavi et al. (2020) can be loaded. - - .. document_args:: seisbench.models EQTransformer - - :param in_channels: Number of input channels, by default 3. - :param in_samples: Number of input samples per channel, by default 6000. - The model expects input shape (in_channels, in_samples) - :param classes: Number of output classes, by default 2. The detection channel is not counted. - :param phases: Phase hints for the classes, by default "PS". Can be None. - :param res_cnn_blocks: Number of residual convolutional blocks - :param lstm_blocks: Number of LSTM blocks - :param drop_rate: Dropout rate - :param original_compatible: If True, uses a few custom layers for binary compatibility with original model - from Mousavi et al. (2020). - This option defaults to False. - It is usually recommended to stick to the default value, as the custom layers show - slightly worse performance than the PyTorch builtins. - The exception is when loading the original weights using :py:func:`from_pretrained`. - :param norm: Data normalization strategy, either "peak" or "std". - :param kwargs: Keyword arguments passed to the constructor of :py:class:`WaveformModel`. - """ - - - - def __init__( - self, - in_channels=3, - in_samples=6000, - classes=2, - phases="PS", - lstm_blocks=3, - drop_rate=0.1, - original_compatible=False, - sampling_rate=100, - norm="std", - **kwargs, - ): - citation = ( - "Mousavi, S.M., Ellsworth, W.L., Zhu, W., Chuang, L, Y., and Beroza, G, C. " - "Earthquake transformer—an attentive deep-learning model for simultaneous earthquake " - "detection and phase picking. Nat Commun 11, 3952 (2020). " - "https://doi.org/10.1038/s41467-020-17591-w" - ) - - # PickBlue options - for option in ("norm_amp_per_comp", "norm_detrend"): - if option in kwargs: - setattr(self, option, kwargs[option]) - del kwargs[option] - else: - setattr(self, option, False) - - # Blinding defines how many samples at beginning and end of the prediction should be ignored - # This is usually required to mitigate prediction problems from training properties, e.g., - # if all picks in the training fall between seconds 5 and 55. - super().__init__() - - self.in_channels = in_channels - self.classes = classes - self.lstm_blocks = lstm_blocks - self.drop_rate = drop_rate - self.norm = norm - self.in_samples = in_samples - - # Add options for conservative and the true original - see https://github.com/seisbench/seisbench/issues/96#issuecomment-1155158224 - if original_compatible: - warnings.warn( - "Using the non-conservative 'original' model, set `original_compatible='conservative' to use the more conservative model" - ) - original_compatible = "non-conservative" - - if original_compatible: - eps = 1e-7 # See Issue #96 - original models use tensorflow default epsilon of 1e-7 - else: - eps = 1e-5 - self.original_compatible = original_compatible - - if original_compatible and in_samples != 6000: - raise ValueError("original_compatible=True requires in_samples=6000.") - - self._phases = phases - if phases is not None and len(phases) != classes: - raise ValueError( - f"Number of classes ({classes}) does not match number of phases ({len(phases)})." - ) - - # Parameters from EQTransformer repository - self.filters = [ - 8, - 16, - 16, - 32, - 32, - 64, - 64, - ] # Number of filters for the convolutions - self.kernel_sizes = [11, 9, 7, 7, 5, 5, 3] # Kernel sizes for the convolutions - self.res_cnn_kernels = [3, 3, 3, 3, 2, 3, 2] - - # TODO: Add regularizers when training model - # kernel_regularizer=keras.regularizers.l2(1e-6), - # bias_regularizer=keras.regularizers.l1(1e-4), - - # Encoder stack - self.encoder = Encoder( - input_channels=self.in_channels, - filters=self.filters, - kernel_sizes=self.kernel_sizes, - in_samples=self.in_samples, - ) - - # Res CNN Stack - self.res_cnn_stack = ResCNNStack( - kernel_sizes=self.res_cnn_kernels, - filters=self.filters[-1], - drop_rate=self.drop_rate, - ) - - # BiLSTM stack - self.bi_lstm_stack = BiLSTMStack( - blocks=self.lstm_blocks, - input_size=self.filters[-1], - drop_rate=self.drop_rate, - original_compatible=original_compatible, - ) - - # Global attention - two transformers - self.transformer_d0 = Transformer( - input_size=16, drop_rate=self.drop_rate, eps=eps - ) - self.transformer_d = Transformer( - input_size=16, drop_rate=self.drop_rate, eps=eps - ) - - # Detection decoder and final Conv - self.decoder_d = Decoder( - input_channels=16, - filters=self.filters[::-1], - kernel_sizes=self.kernel_sizes[::-1], - out_samples=in_samples, - original_compatible=original_compatible, - ) - self.conv_d = nn.Conv1d( - in_channels=self.filters[0], out_channels=1, kernel_size=11, padding=5 - ) - - # Picking branches - self.pick_lstms = [] - self.pick_attentions = [] - self.pick_decoders = [] - self.pick_convs = [] - self.dropout = nn.Dropout(drop_rate) - - for _ in range(self.classes): - if original_compatible == "conservative": - # The non-conservative model uses a sigmoid activiation as handled by the base nn.LSTM - lstm = CustomLSTM(ActivationLSTMCell, 16, 16, bidirectional=False) - else: - lstm = nn.LSTM(16, 16, bidirectional=False) - self.pick_lstms.append(lstm) - - attention = SeqSelfAttention(input_size=16, attention_width=3, eps=eps) - self.pick_attentions.append(attention) - - decoder = Decoder( - input_channels=16, - filters=self.filters[::-1], - kernel_sizes=self.kernel_sizes[::-1], - out_samples=in_samples, - original_compatible=original_compatible, - ) - self.pick_decoders.append(decoder) - - conv = nn.Conv1d( - in_channels=self.filters[0], out_channels=1, kernel_size=11, padding=5 - ) - self.pick_convs.append(conv) - - self.pick_lstms = nn.ModuleList(self.pick_lstms) - self.pick_attentions = nn.ModuleList(self.pick_attentions) - self.pick_decoders = nn.ModuleList(self.pick_decoders) - self.pick_convs = nn.ModuleList(self.pick_convs) - - def forward(self, x): - assert x.ndim == 3 - assert x.shape[1:] == (self.in_channels, self.in_samples) - - # Shared encoder part - x = self.encoder(x) - x = self.res_cnn_stack(x) - x = self.bi_lstm_stack(x) - x, _ = self.transformer_d0(x) - x, _ = self.transformer_d(x) - - # Detection part - detection = self.decoder_d(x) - - detection = torch.sigmoid(self.conv_d(detection)) - detection = torch.squeeze(detection, dim=1) # Remove channel dimension - - outputs = [detection] - - # Pick parts - for lstm, attention, decoder, conv in zip( - self.pick_lstms, self.pick_attentions, self.pick_decoders, self.pick_convs - ): - px = x.permute( - 2, 0, 1 - ) # From batch, channels, sequence to sequence, batch, channels - px = lstm(px)[0] - px = self.dropout(px) - px = px.permute( - 1, 2, 0 - ) # From sequence, batch, channels to batch, channels, sequence - px, _ = attention(px) - px = decoder(px) - - pred = torch.sigmoid(conv(px)) - pred = torch.squeeze(pred, dim=1) # Remove channel dimension - - outputs.append(pred) - - return tuple(outputs) - - - - -class Encoder(nn.Module): - """ - Encoder stack - """ - - def __init__(self, input_channels, filters, kernel_sizes, in_samples): - super().__init__() - - convs = [] - pools = [] - self.paddings = [] - for in_channels, out_channels, kernel_size in zip( - [input_channels] + filters[:-1], filters, kernel_sizes - ): - convs.append( - nn.Conv1d( - in_channels, out_channels, kernel_size, padding=kernel_size // 2 - ) - ) - - # To be consistent with the behaviour in tensorflow, - # padding needs to be added for odd numbers of input_samples - padding = in_samples % 2 - - # Padding for MaxPool1d needs to be handled manually to conform with tf padding - self.paddings.append(padding) - pools.append(nn.MaxPool1d(2, padding=0)) - in_samples = (in_samples + padding) // 2 - - self.convs = nn.ModuleList(convs) - self.pools = nn.ModuleList(pools) - #self.paddings = nn.ModuleList(self.paddings) - - def forward(self, x): - # Use zip for convs and pools, access paddings by index - for i, (conv, pool) in enumerate(zip(self.convs, self.pools)): - padding = self.paddings[i] # Access padding by index - - x = torch.relu(conv(x)) # Apply convolution followed by ReLU - if padding != 0: - # Only pad right, use -1e10 as negative infinity - x = F.pad(x, (0, padding), "constant", -1000000.0) # Apply padding to the right - x = pool(x) # Apply pooling - - return x - - -class Decoder(nn.Module): - def __init__( - self, - input_channels, - filters, - kernel_sizes, - out_samples, - original_compatible=False, - ): - super().__init__() - - self.upsample = nn.Upsample(scale_factor=2, mode="nearest") - self.original_compatible = original_compatible - - # We need to trim off the final sample sometimes to get to the right number of output samples - self.crops = [] - current_samples = out_samples - for i, _ in enumerate(filters): - padding = current_samples % 2 - current_samples = (current_samples + padding) // 2 - if padding == 1: - self.crops.append(len(filters) - 1 - i) - - convs = [] - for in_channels, out_channels, kernel_size in zip( - [input_channels] + filters[:-1], filters, kernel_sizes - ): - convs.append( - nn.Conv1d( - in_channels, out_channels, kernel_size, padding=kernel_size // 2 - ) - ) - - self.convs = nn.ModuleList(convs) - - def forward(self, x): - for i, conv in enumerate(self.convs): - x = self.upsample(x) - - if self.original_compatible: - if i == 3: - x = x[:, :, 1:-1] - else: - if i in self.crops: - x = x[:, :, :-1] - - x = F.relu(conv(x)) - - return x - - -class ResCNNStack(nn.Module): - def __init__(self, kernel_sizes, filters, drop_rate): - super().__init__() - - members = [] - for ker in kernel_sizes: - members.append(ResCNNBlock(filters, ker, drop_rate)) - - self.members = nn.ModuleList(members) - - def forward(self, x): - for member in self.members: - x = member(x) - - return x - - -class ResCNNBlock(nn.Module): - def __init__(self, filters, ker, drop_rate): - super().__init__() - - self.manual_padding = False - if ker == 3: - padding = 1 - else: - # ker == 2 - # Manual padding emulate the padding in tensorflow - self.manual_padding = True - padding = 0 - - self.dropout = SpatialDropout1d(drop_rate) - - self.norm1 = nn.BatchNorm1d(filters, eps=1e-3) - self.conv1 = nn.Conv1d(filters, filters, ker, padding=padding) - - self.norm2 = nn.BatchNorm1d(filters, eps=1e-3) - self.conv2 = nn.Conv1d(filters, filters, ker, padding=padding) - - def forward(self, x): - y = self.norm1(x) - y = F.relu(y) - y = self.dropout(y) - if self.manual_padding: - y = F.pad(y, (0, 1), "constant", 0.0) - y = self.conv1(y) - - y = self.norm2(y) - y = F.relu(y) - y = self.dropout(y) - if self.manual_padding: - y = F.pad(y, (0, 1), "constant", 0.0) - y = self.conv2(y) - - return x + y - - -class BiLSTMStack(nn.Module): - def __init__( - self, blocks, input_size, drop_rate, hidden_size=16, original_compatible=False - ): - super().__init__() - - # First LSTM has a different input size as the subsequent ones - self.members = nn.ModuleList( - [ - BiLSTMBlock( - input_size, - hidden_size, - drop_rate, - original_compatible=original_compatible, - ) - ] - + [ - BiLSTMBlock( - hidden_size, - hidden_size, - drop_rate, - original_compatible=original_compatible, - ) - for _ in range(blocks - 1) - ] - ) - - def forward(self, x): - for member in self.members: - x = member(x) - return x - - -class BiLSTMBlock(nn.Module): - def __init__(self, input_size, hidden_size, drop_rate, original_compatible=False): - super().__init__() - - if original_compatible == "conservative": - # The non-conservative model uses a sigmoid activiation as handled by the base nn.LSTM - self.lstm = CustomLSTM(ActivationLSTMCell, input_size, hidden_size) - elif original_compatible == "non-conservative": - self.lstm = CustomLSTM( - ActivationLSTMCell, - input_size, - hidden_size, - gate_activation=torch.sigmoid, - ) - else: - self.lstm = nn.LSTM(input_size, hidden_size, bidirectional=True) - self.dropout = nn.Dropout(drop_rate) - self.conv = nn.Conv1d(2 * hidden_size, hidden_size, 1) - self.norm = nn.BatchNorm1d(hidden_size, eps=1e-3) - - def forward(self, x): - x = x.permute( - 2, 0, 1 - ) # From batch, channels, sequence to sequence, batch, channels - x = self.lstm(x)[0] - x = self.dropout(x) - x = x.permute( - 1, 2, 0 - ) # From sequence, batch, channels to batch, channels, sequence - x = self.conv(x) - x = self.norm(x) - return x - - -class Transformer(nn.Module): - def __init__(self, input_size, drop_rate, attention_width=None, eps=1e-5): - super().__init__() - - self.attention = SeqSelfAttention( - input_size, attention_width=attention_width, eps=eps - ) - self.norm1 = LayerNormalization(input_size) - self.ff = FeedForward(input_size, drop_rate) - self.norm2 = LayerNormalization(input_size) - - def forward(self, x): - y, weight = self.attention(x) - y = x + y - y = self.norm1(y) - y2 = self.ff(y) - y2 = y + y2 - y2 = self.norm2(y2) - - return y2, weight - - -class SeqSelfAttention(nn.Module): - """ - Additive self attention - """ - - def __init__(self, input_size, units=32, attention_width=None, eps=1e-5): - super().__init__() - self.attention_width = attention_width - - self.Wx = nn.Parameter(uniform(-0.02, 0.02, input_size, units)) - self.Wt = nn.Parameter(uniform(-0.02, 0.02, input_size, units)) - self.bh = nn.Parameter(torch.zeros(units)) - - self.Wa = nn.Parameter(uniform(-0.02, 0.02, units, 1)) - self.ba = nn.Parameter(torch.zeros(1)) - - self.eps = eps - - def forward(self, x): - # x.shape == (batch, channels, time) - - x = x.permute(0, 2, 1) # to (batch, time, channels) - - q = torch.unsqueeze( - torch.matmul(x, self.Wt), 2 - ) # Shape (batch, time, 1, channels) - k = torch.unsqueeze( - torch.matmul(x, self.Wx), 1 - ) # Shape (batch, 1, time, channels) - - h = torch.tanh(q + k + self.bh) - - # Emissions - e = torch.squeeze( - torch.matmul(h, self.Wa) + self.ba, -1 - ) # Shape (batch, time, time) - - # This is essentially softmax with an additional attention component. - e = ( - e - torch.max(e, dim=-1, keepdim=True).values - ) # In versions <= 0.2.1 e was incorrectly normalized by max(x) - e = torch.exp(e) - if self.attention_width is not None: - lower = ( - torch.arange(0, e.shape[1], device=e.device) - self.attention_width // 2 - ) - upper = lower + self.attention_width - indices = torch.unsqueeze(torch.arange(0, e.shape[1], device=e.device), 1) - mask = torch.logical_and(lower <= indices, indices < upper) - e = torch.where(mask, e, torch.zeros_like(e)) - - a = e / (torch.sum(e, dim=-1, keepdim=True) + self.eps) - - v = torch.matmul(a, x) - - v = v.permute(0, 2, 1) # to (batch, channels, time) - - return v, a - - -def uniform(a, b, *args): - return a + (b - a) * torch.rand(*args) - - -class LayerNormalization(nn.Module): - def __init__(self, filters, eps=1e-14): - super().__init__() - - gamma = torch.ones(filters, 1) - self.gamma = nn.Parameter(gamma) - beta = torch.zeros(filters, 1) - self.beta = nn.Parameter(beta) - self.eps = eps - - def forward(self, x): - mean = torch.mean(x, 1, keepdim=True) - var = torch.mean((x - mean) ** 2, 1, keepdim=True) + self.eps - std = torch.sqrt(var) - outputs = (x - mean) / std - - outputs = outputs * self.gamma - outputs = outputs + self.beta - - return outputs - - -class FeedForward(nn.Module): - def __init__(self, io_size, drop_rate, hidden_size=128): - super().__init__() - - self.lin1 = nn.Linear(io_size, hidden_size) - self.lin2 = nn.Linear(hidden_size, io_size) - self.dropout = nn.Dropout(drop_rate) - - def forward(self, x): - x = x.permute(0, 2, 1) # To (batch, time, channel) - x = F.relu(self.lin1(x)) - x = self.dropout(x) - x = self.lin2(x) - x = x.permute(0, 2, 1) # To (batch, channel, time) - - return x - - -class SpatialDropout1d(nn.Module): - def __init__(self, drop_rate): - super().__init__() - - self.drop_rate = drop_rate - self.dropout = nn.Dropout2d(drop_rate) - - def forward(self, x): - x = x.unsqueeze(dim=-1) # Add fake dimension - x = self.dropout(x) - x = x.squeeze(dim=-1) # Remove fake dimension - return x - +from models.EQTransformer import EQTransformer +from onnx_picker_base import OnnxSlidingWindowPicker +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +pretrained = EQTransformer.from_pretrained("stead").to(device) +pretrained_state = {k: v.cpu() for k, v in pretrained.state_dict().items()} -class Picker(EQTransformerForPicker): +class Picker(OnnxSlidingWindowPicker): def __init__(self): - super().__init__() - - def forward(self, x): - device = x.device - with torch.no_grad(): - #print("数据维度", x.shape) - T, C = x.shape - seqlen = 6000 - batchstride = seqlen - 3000 - batchlen = torch.ceil(torch.tensor(T / batchstride).to(device)) - idx = torch.arange(0, seqlen, 1, device=device).unsqueeze(0) + torch.arange(0, batchlen, 1, device=device).unsqueeze(1) * batchstride - idx = idx.clamp(min=0, max=T-1).long() - x = x.to(device) - wave = x[idx, :] - wave = wave.permute(0, 2, 1) - wave -= torch.mean(wave, dim=2, keepdim=True) - #max, maxidx = torch.max(torch.abs(wave), dim=2, keepdim=True) - #max, maxidx = torch.max(max, dim=1, keepdim=True) - max = torch.std(wave, dim=2, keepdim=True) - wave /= (max + 1e-6) - x = wave - x = self.encoder(x) - x = self.res_cnn_stack(x) - x = self.bi_lstm_stack(x) - x, _ = self.transformer_d0(x) - x, _ = self.transformer_d(x) - - # Detection part - detection = self.decoder_d(x) - - detection = torch.sigmoid(self.conv_d(detection)) - detection = torch.squeeze(detection, dim=1) # Remove channel dimension - - outputs = [detection] - - # Pick parts - for lstm, attention, decoder, conv in zip( - self.pick_lstms, self.pick_attentions, self.pick_decoders, self.pick_convs - ): - px = x.permute( - 2, 0, 1 - ) # From batch, channels, sequence to sequence, batch, channels - px = lstm(px)[0] - px = self.dropout(px) - px = px.permute( - 1, 2, 0 - ) # From sequence, batch, channels to batch, channels, sequence - px, _ = attention(px) - px = decoder(px) + super().__init__(EQTransformer, state_dict=pretrained_state, seqlen=6000, overlap=3000) - pred = torch.sigmoid(conv(px)) - pred = torch.squeeze(pred, dim=1) # Remove channel dimension - outputs.append(pred) - oc = torch.stack(outputs, dim=1) - B, C, T = oc.shape - tgrid = torch.arange(0, T, 1, device=device).unsqueeze(0) * 1 + torch.arange(0, batchlen, 1, device=device).unsqueeze(1) * batchstride - oc = oc.permute(0, 2, 1).reshape(-1, C) - #oc = oc[:, [2, 0, 1]] - ot = tgrid.squeeze() - ot = ot.reshape(-1) - return oc, ot - - - - -# Define the device (use 'cuda' if you have a GPU, else 'cpu') -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - -# Load the pre-trained EqTransformer model from SeisBench (you can replace with your own model path) -model2 = EQTransformer.from_pretrained('stead').to(device) - -model = Picker() -model.load_state_dict(model2.state_dict()) +model = Picker() model.eval() - input_names = ["wave"] output_names = ["prob", "time"] -#x = torch.randn([10, 3, 6144, 1]) x = torch.randn([500000, 3]) -torch.onnx.export(model, x, -"pickers/eqtransformer.stead.onnx", verbose=True, -dynamic_axes={"wave":{0:"batch"}, "prob":{0:"batch"}, "time":{0:"batch"}}, -input_names=input_names, output_names=output_names, opset_version=11) +torch.onnx.export( + model, + x, + "pickers/eqtransformer.stead.onnx", + verbose=True, + dynamic_axes={"wave": {0: "batch"}, "prob": {0: "batch"}, "time": {0: "batch"}}, + input_names=input_names, + output_names=output_names, + opset_version=11, +) diff --git a/makeonnx.lppn.py b/makeonnx.lppn.py index 8a784ea..c3f47d4 100755 --- a/makeonnx.lppn.py +++ b/makeonnx.lppn.py @@ -1,61 +1,59 @@ -from models.LPPNM import Model -import torch +import torch +from onnx_picker_base import OnnxSlidingWindowPicker -modelname = "lppnm" -if modelname == "lppnt": +MODEL_NAME = "lppnm" + +if MODEL_NAME == "lppnt": from models.LPPNT import Model -elif modelname == "lppnm": - from models.LPPNM import Model -elif modelname == "lppnl": + CKPT_PATH = "ckpt/china.lppnt.pt" +elif MODEL_NAME == "lppnl": from models.LPPNL import Model -class Picker(Model): - def __init__(self, n_stride=8): - super().__init__() - def forward(self, x): + CKPT_PATH = "ckpt/china.lppnl.pt" +else: + from models.LPPNM import Model + CKPT_PATH = "ckpt/china.lppnm.pt" + + +class Picker(OnnxSlidingWindowPicker): + def __init__(self): + super().__init__(Model, ckpt_path=CKPT_PATH) self.n_stride = 8 - device = x.device + + def forward(self, x): + device = x.device with torch.no_grad(): - #print("数据维度", x.shape) - T, C = x.shape - seqlen = 6144 - batchstride = 6144 - 256 - batchlen = torch.ceil(torch.tensor(T / batchstride).to(device)) - idx = torch.arange(0, seqlen, 1, device=device).unsqueeze(0) + torch.arange(0, batchlen, 1, device=device).unsqueeze(1) * batchstride - idx = idx.clamp(min=0, max=T-1).long() - x = x.to(device) - wave = x[idx, :] - wave = wave.permute(0, 2, 1) - wave -= torch.mean(wave, dim=2, keepdim=True) - max1, max1idx = torch.max(torch.abs(wave), dim=2, keepdim=True) - max2, max2idx = torch.max(max1, dim=2, keepdim=True) - wave /= (max2 + 1e-6) + wave, batchlen = self.window_and_normalize(x, device) wave = wave.unsqueeze(2) - x1 = self.layers(wave) - x2 = self.class_encoder(x1) - x = torch.cat([x1, x2], dim=1) - out_class = self.cl(x).squeeze(dim=2) - out_time = self.tm(x) - out_time = out_time.sigmoid().squeeze() * self.n_stride - - oc = out_class.squeeze() - ot = out_time.squeeze() - #print(oc.shape, ot.shape) - B, C, T = oc.shape - oc = oc.softmax(dim=1) - tgrid = torch.arange(0, T, 1, device=device).unsqueeze(0) * self.n_stride + torch.arange(0, batchlen, 1, device=device).unsqueeze(1) * batchstride - oc = oc.permute(0, 2, 1).reshape(-1, C) - ot += tgrid.squeeze() - ot = ot.reshape(-1) - return oc, ot - -model = Picker() -model.load_state_dict(torch.load(f"ckpt/china.{modelname}.pt", map_location="cpu")) + + x1 = self.model.layers(wave) + x2 = self.model.class_encoder(x1) + features = torch.cat([x1, x2], dim=1) + out_class = self.model.cl(features).squeeze(dim=2) + out_time = self.model.tm(features).sigmoid().squeeze() * self.n_stride + + oc = out_class.softmax(dim=1) + B, C, T = oc.shape + tgrid = ( + torch.arange(0, T, 1, device=device).unsqueeze(0) * self.n_stride + + torch.arange(0, batchlen, 1, device=device).unsqueeze(1) * self.batchstride + ) + ot = (out_time + tgrid).reshape(-1) + oc = oc.permute(0, 2, 1).reshape(-1, C) + return oc, ot + + +model = Picker() model.eval() input_names = ["wave"] output_names = ["prob", "time"] -#x = torch.randn([10, 3, 6144, 1]) x = torch.randn([500000, 3]) -torch.onnx.export(model, x, -f"pickers/{modelname}.onnx", verbose=True, -dynamic_axes={"wave":{0:"batch"}, "prob":{0:"batch"}, "time":{0:"batch"}}, -input_names=input_names, output_names=output_names, opset_version=11) +torch.onnx.export( + model, + x, + f"pickers/{MODEL_NAME}.onnx", + verbose=True, + dynamic_axes={"wave": {0: "batch"}, "prob": {0: "batch"}, "time": {0: "batch"}}, + input_names=input_names, + output_names=output_names, + opset_version=11, +) diff --git a/makeonnx.pnsn.py b/makeonnx.pnsn.py index 15647ba..e3f8614 100755 --- a/makeonnx.pnsn.py +++ b/makeonnx.pnsn.py @@ -1,44 +1,25 @@ -import torch -import torch.nn as nn -from models.BRNNPNSN import BRNN -class Picker(BRNN): +import torch +from models.BRNNPNSN import BRNN +from onnx_picker_base import OnnxSlidingWindowPicker + + +class Picker(OnnxSlidingWindowPicker): def __init__(self): - super().__init__() - self.n_stride = 1 - def forward(self, x): - device = x.device - with torch.no_grad(): - #print("数据维度", x.shape) - T, C = x.shape - seqlen = 10240 - batchstride = seqlen - 512 - batchlen = torch.ceil(torch.tensor(T / batchstride).to(device)) - idx = torch.arange(0, seqlen, 1, device=device).unsqueeze(0) + torch.arange(0, batchlen, 1, device=device).unsqueeze(1) * batchstride - idx = idx.clamp(min=0, max=T-1).long() - x = x.to(device) - wave = x[idx, :] - wave = wave.permute(0, 2, 1) - wave -= torch.mean(wave, dim=2, keepdim=True) - max = torch.std(wave, dim=2, keepdim=True) - wave /= (max + 1e-6) - x = self.encoder(wave) - e = self.rnns(x) # 波形特征 - y = self.decoder(e) # 输出概率 - oc = y.softmax(dim=1) - B, C, T = oc.shape - tgrid = torch.arange(0, T, 1, device=device).unsqueeze(0) * self.n_stride + torch.arange(0, batchlen, 1, device=device).unsqueeze(1) * batchstride - oc = oc.permute(0, 2, 1).reshape(-1, C) - ot = tgrid.squeeze() - ot = ot.reshape(-1) - return oc, ot -model = Picker() + super().__init__(BRNN, ckpt_path="ckpt/china.rnn.pnsn.pt", seqlen=10240, overlap=512) + + +model = Picker() model.eval() -model.load_state_dict(torch.load("ckpt/china.rnn.pnsn.pt", map_location="cpu")) input_names = ["wave"] output_names = ["prob", "time"] -#x = torch.randn([10, 3, 6144, 1]) x = torch.randn([500000, 3]) -torch.onnx.export(model, x, -"pickers/rnn.pnsn.onnx", verbose=True, -dynamic_axes={"wave":{0:"batch"}, "prob":{0:"batch"}, "time":{0:"batch"}}, -input_names=input_names, output_names=output_names, opset_version=11) \ No newline at end of file +torch.onnx.export( + model, + x, + "pickers/rnn.pnsn.onnx", + verbose=True, + dynamic_axes={"wave": {0: "batch"}, "prob": {0: "batch"}, "time": {0: "batch"}}, + input_names=input_names, + output_names=output_names, + opset_version=11, +) diff --git a/makeonnx.rnn.py b/makeonnx.rnn.py index 9bc3ab3..ed7cddf 100755 --- a/makeonnx.rnn.py +++ b/makeonnx.rnn.py @@ -1,44 +1,25 @@ -import torch -import torch.nn as nn -from models.BRNN import BRNN -class Picker(BRNN): +import torch +from models.BRNN import BRNN +from onnx_picker_base import OnnxSlidingWindowPicker + + +class Picker(OnnxSlidingWindowPicker): def __init__(self): - super().__init__() - self.n_stride = 1 - def forward(self, x): - device = x.device - with torch.no_grad(): - #print("数据维度", x.shape) - T, C = x.shape - seqlen = 6144 - batchstride = 6144 - 256 - batchlen = torch.ceil(torch.tensor(T / batchstride).to(device)) - idx = torch.arange(0, seqlen, 1, device=device).unsqueeze(0) + torch.arange(0, batchlen, 1, device=device).unsqueeze(1) * batchstride - idx = idx.clamp(min=0, max=T-1).long() - x = x.to(device) - wave = x[idx, :] - wave = wave.permute(0, 2, 1) - wave -= torch.mean(wave, dim=2, keepdim=True) - max, maxidx = torch.max(torch.abs(wave), dim=2, keepdim=True) - wave /= (max + 1e-6) - x = self.encoder(wave) - e = self.rnns(x) # 波形特征 - y = self.decoder(e) # 输出概率 - oc = y.softmax(dim=1) - B, C, T = oc.shape - tgrid = torch.arange(0, T, 1, device=device).unsqueeze(0) * self.n_stride + torch.arange(0, batchlen, 1, device=device).unsqueeze(1) * batchstride - oc = oc.permute(0, 2, 1).reshape(-1, C) - ot = tgrid.squeeze() - ot = ot.reshape(-1) - return oc, ot -model = Picker() + super().__init__(BRNN, ckpt_path="ckpt/china.rnn.pt") + + +model = Picker() model.eval() -model.load_state_dict(torch.load("ckpt/china.rnn.pt", map_location="cpu")) input_names = ["wave"] output_names = ["prob", "time"] -#x = torch.randn([10, 3, 6144, 1]) x = torch.randn([500000, 3]) -torch.onnx.export(model, x, -"pickers/rnn.onnx", verbose=True, -dynamic_axes={"wave":{0:"batch"}, "prob":{0:"batch"}, "time":{0:"batch"}}, -input_names=input_names, output_names=output_names, opset_version=11) \ No newline at end of file +torch.onnx.export( + model, + x, + "pickers/rnn.onnx", + verbose=True, + dynamic_axes={"wave": {0: "batch"}, "prob": {0: "batch"}, "time": {0: "batch"}}, + input_names=input_names, + output_names=output_names, + opset_version=11, +) diff --git a/makeonnx.skynet.multiphase.py b/makeonnx.skynet.multiphase.py index fdb6763..e20b27f 100644 --- a/makeonnx.skynet.multiphase.py +++ b/makeonnx.skynet.multiphase.py @@ -3,6 +3,7 @@ import numpy as np import torch import torch.nn as nn +from onnx_picker_base import OnnxSlidingWindowPicker class Skynet(nn.Module): """ @@ -180,89 +181,25 @@ def forward(self, X): from models.UNet import PhaseNetLight import torch -class Picker(Skynet): + + +class Picker(OnnxSlidingWindowPicker): def __init__(self): - super().__init__() - - def forward(self, x): - device = x.device - with torch.no_grad(): - #print("数据维度", x.shape) - T, C = x.shape - seqlen = 30000 - batchstride = seqlen - 30000 // 2 - batchlen = torch.ceil(torch.tensor(T / batchstride).to(device)) - idx = torch.arange(0, seqlen, 1, device=device).unsqueeze(0) + torch.arange(0, batchlen, 1, device=device).unsqueeze(1) * batchstride - idx = idx.clamp(min=0, max=T-1).long() - x = x.to(device) - wave = x[idx, :] - wave = wave.permute(0, 2, 1) - wave -= torch.mean(wave, dim=2, keepdim=True) - max, maxidx = torch.max(torch.abs(wave), dim=2, keepdim=True) - max, maxidx = torch.max(max, dim=1, keepdim=True) - #max = torch.std(wave, dim=2, keepdim=True) - wave /= (max + 1e-6) - X1 = torch.relu(self.bn1(self.conv1(wave))) - X2 = torch.relu(self.bn2(self.conv2(X1))) - X3 = torch.relu(self.bn3(self.conv3(X2))) - X4 = torch.relu(self.bn4(self.conv4(X3))) - X5 = torch.relu(self.bn5(self.conv5(X4))) - X6 = torch.relu(self.bn6(self.conv6(X5))) - X7 = torch.relu(self.bn7(self.conv7(X6))) - X8 = torch.relu(self.bn8(self.conv8(X7))) - X9 = torch.relu(self.bn9(self.conv9(X8))) - X10 = torch.relu(self.bn10(self.conv10(X9))) - # extra from original UNet - X10_a = torch.relu(self.bn11(self.conv11(X10))) - X10_b = torch.relu(self.bn12(self.conv12(X10_a))) - X10_c = torch.relu(self.bnd0(self.dconv0(X10_b))) - X10_c = torch.cat( - ( - X10_c, - torch.zeros((X10_c.shape[0], X10_c.shape[1], 1), device=X10_c.device), - ), - dim=-1, - ) - X10_c = torch.cat((X10, X10_c), dim=1) - X10_d = torch.relu(self.bnd01(self.dconv01(X10_c))) - X11 = torch.relu(self.bnd1(self.dconv1(X10_d))) - X12 = torch.cat((X11, X8), dim=1) - X12 = torch.relu(self.bnd2(self.dconv2(X12))) - X13 = torch.relu(self.bnd3(self.dconv3(X12))) - X14 = torch.relu(self.bnd4(self.dconv4(torch.cat((X13, X6), dim=1)))) - X15 = torch.relu(self.bnd5(self.dconv5(X14))) - X15 = torch.cat( - (X15, torch.zeros((X15.shape[0], X15.shape[1], 1), device=X15.device)), - dim=2, - ) - X16 = torch.relu(self.bnd6(self.dconv6(torch.cat((X15, X4), dim=1)))) - X17 = torch.relu(self.bnd7(self.dconv7(X16))) - X17 = torch.cat( - (X17, torch.zeros((X17.shape[0], X17.shape[1], 1), device=X17.device)), - dim=2, - ) - X18 = torch.relu(self.bnd8(self.dconv8(torch.cat((X17, X2), dim=1)))) - X19 = self.dconv9(X18) + super().__init__(Skynet, ckpt_path="skynet/skynet_models/seisbench_skynet_multiphase.pt", seqlen=30000, overlap=15000) + - oc = self.softmax(X19) - B, C, T = oc.shape - tgrid = torch.arange(0, T, 1, device=device).unsqueeze(0) * 1 + torch.arange(0, batchlen, 1, device=device).unsqueeze(1) * batchstride - oc = oc.permute(0, 2, 1).reshape(-1, C) - oc = oc[:, [4, 1, 3, 0, 2]] - ot = tgrid.squeeze() - ot = ot.reshape(-1) - return oc, ot -model = Picker() +model = Picker() model.eval() -ckpt = torch.load("skynet/skynet_models/multiphase_skynet.pt", weights_only=False, map_location="cpu") -#state = ckpt.state_dict() -print(ckpt) -model.load_state_dict(ckpt) input_names = ["wave"] output_names = ["prob", "time"] -#x = torch.randn([10, 3, 6144, 1]) x = torch.randn([500000, 3]) -torch.onnx.export(model, x, -"pickers/skynet.multiphase.onnx", verbose=True, -dynamic_axes={"wave":{0:"batch"}, "prob":{0:"batch"}, "time":{0:"batch"}}, -input_names=input_names, output_names=output_names, opset_version=11) +torch.onnx.export( + model, + x, + "pickers/skynet.multiphase.onnx", + verbose=True, + dynamic_axes={"wave": {0: "batch"}, "prob": {0: "batch"}, "time": {0: "batch"}}, + input_names=input_names, + output_names=output_names, + opset_version=11, +) diff --git a/makeonnx.skynet.py b/makeonnx.skynet.py index ed40b16..9694f64 100644 --- a/makeonnx.skynet.py +++ b/makeonnx.skynet.py @@ -3,6 +3,7 @@ import numpy as np import torch import torch.nn as nn +from onnx_picker_base import OnnxSlidingWindowPicker class Skynet(nn.Module): """ @@ -180,89 +181,25 @@ def forward(self, X): from models.UNet import PhaseNetLight import torch -class Picker(Skynet): + + +class Picker(OnnxSlidingWindowPicker): def __init__(self): - super().__init__() - - def forward(self, x): - device = x.device - with torch.no_grad(): - #print("数据维度", x.shape) - T, C = x.shape - seqlen = 30000 - batchstride = seqlen - 30000 // 2 - batchlen = torch.ceil(torch.tensor(T / batchstride).to(device)) - idx = torch.arange(0, seqlen, 1, device=device).unsqueeze(0) + torch.arange(0, batchlen, 1, device=device).unsqueeze(1) * batchstride - idx = idx.clamp(min=0, max=T-1).long() - x = x.to(device) - wave = x[idx, :] - wave = wave.permute(0, 2, 1) - wave -= torch.mean(wave, dim=2, keepdim=True) - max, maxidx = torch.max(torch.abs(wave), dim=2, keepdim=True) - max, maxidx = torch.max(max, dim=1, keepdim=True) - #max = torch.std(wave, dim=2, keepdim=True) - wave /= (max + 1e-6) - X1 = torch.relu(self.bn1(self.conv1(wave))) - X2 = torch.relu(self.bn2(self.conv2(X1))) - X3 = torch.relu(self.bn3(self.conv3(X2))) - X4 = torch.relu(self.bn4(self.conv4(X3))) - X5 = torch.relu(self.bn5(self.conv5(X4))) - X6 = torch.relu(self.bn6(self.conv6(X5))) - X7 = torch.relu(self.bn7(self.conv7(X6))) - X8 = torch.relu(self.bn8(self.conv8(X7))) - X9 = torch.relu(self.bn9(self.conv9(X8))) - X10 = torch.relu(self.bn10(self.conv10(X9))) - # extra from original UNet - X10_a = torch.relu(self.bn11(self.conv11(X10))) - X10_b = torch.relu(self.bn12(self.conv12(X10_a))) - X10_c = torch.relu(self.bnd0(self.dconv0(X10_b))) - X10_c = torch.cat( - ( - X10_c, - torch.zeros((X10_c.shape[0], X10_c.shape[1], 1), device=X10_c.device), - ), - dim=-1, - ) - X10_c = torch.cat((X10, X10_c), dim=1) - X10_d = torch.relu(self.bnd01(self.dconv01(X10_c))) - X11 = torch.relu(self.bnd1(self.dconv1(X10_d))) - X12 = torch.cat((X11, X8), dim=1) - X12 = torch.relu(self.bnd2(self.dconv2(X12))) - X13 = torch.relu(self.bnd3(self.dconv3(X12))) - X14 = torch.relu(self.bnd4(self.dconv4(torch.cat((X13, X6), dim=1)))) - X15 = torch.relu(self.bnd5(self.dconv5(X14))) - X15 = torch.cat( - (X15, torch.zeros((X15.shape[0], X15.shape[1], 1), device=X15.device)), - dim=2, - ) - X16 = torch.relu(self.bnd6(self.dconv6(torch.cat((X15, X4), dim=1)))) - X17 = torch.relu(self.bnd7(self.dconv7(X16))) - X17 = torch.cat( - (X17, torch.zeros((X17.shape[0], X17.shape[1], 1), device=X17.device)), - dim=2, - ) - X18 = torch.relu(self.bnd8(self.dconv8(torch.cat((X17, X2), dim=1)))) - X19 = self.dconv9(X18) + super().__init__(Skynet, ckpt_path="skynet/skynet_models/seisbench_skynet.pt", seqlen=30000, overlap=15000) + - oc = self.softmax(X19) - B, C, T = oc.shape - tgrid = torch.arange(0, T, 1, device=device).unsqueeze(0) * 1 + torch.arange(0, batchlen, 1, device=device).unsqueeze(1) * batchstride - oc = oc.permute(0, 2, 1).reshape(-1, C) - oc = oc[:, [2, 0, 1]] - ot = tgrid.squeeze() - ot = ot.reshape(-1) - return oc, ot -model = Picker() +model = Picker() model.eval() -ckpt = torch.load("skynet/skynet_models/seisbench_skynet.pt", weights_only=False, map_location="cpu") -#state = ckpt.state_dict() -print(ckpt) -model.load_state_dict(ckpt) input_names = ["wave"] output_names = ["prob", "time"] -#x = torch.randn([10, 3, 6144, 1]) x = torch.randn([500000, 3]) -torch.onnx.export(model, x, -"pickers/skynet.onnx", verbose=True, -dynamic_axes={"wave":{0:"batch"}, "prob":{0:"batch"}, "time":{0:"batch"}}, -input_names=input_names, output_names=output_names, opset_version=11) +torch.onnx.export( + model, + x, + "pickers/skynet.onnx", + verbose=True, + dynamic_axes={"wave": {0: "batch"}, "prob": {0: "batch"}, "time": {0: "batch"}}, + input_names=input_names, + output_names=output_names, + opset_version=11, +) diff --git a/makeonnx.unet.py b/makeonnx.unet.py index 654357c..2f5a611 100644 --- a/makeonnx.unet.py +++ b/makeonnx.unet.py @@ -1,57 +1,25 @@ -from models.UNet import PhaseNetLight -import torch +import torch +from models.UNet import UNet +from onnx_picker_base import OnnxSlidingWindowPicker -class Picker(PhaseNetLight): + +class Picker(OnnxSlidingWindowPicker): def __init__(self): - super().__init__() - - def forward(self, x): - device = x.device - with torch.no_grad(): - #print("数据维度", x.shape) - T, C = x.shape - seqlen = 3072 - batchstride = seqlen - 3072 // 2 - batchlen = torch.ceil(torch.tensor(T / batchstride).to(device)) - idx = torch.arange(0, seqlen, 1, device=device).unsqueeze(0) + torch.arange(0, batchlen, 1, device=device).unsqueeze(1) * batchstride - idx = idx.clamp(min=0, max=T-1).long() - x = x.to(device) - wave = x[idx, :] - wave = wave.permute(0, 2, 1) - wave -= torch.mean(wave, dim=2, keepdim=True) - max, maxidx = torch.max(torch.abs(wave), dim=2, keepdim=True) - max = torch.std(wave, dim=2, keepdim=True) - wave /= (max + 1e-6) - x_in = self.activation(self.in_bn(self.inc(wave))) - x1 = self.activation(self.bnd1(self.conv1(x_in))) - x2 = self.activation(self.bnd2(self.conv2(x1))) - x3 = self.activation(self.bnd3(self.conv3(x2))) - x4 = self.activation(self.bnd4(self.conv4(x3))) - - x = torch.cat([self.activation(self.bnu1(self.up1(x4))), x3], dim=1) - x = torch.cat([self.activation(self.bnu2(self.up2(x))), x2], dim=1) - x = torch.cat([self.activation(self.bnu3(self.up3(x))), x1], dim=1) - x = torch.cat([self.activation(self.bnu4(self.up4(x))), x_in], dim=1) - #print(x.shape) - x = self.out(x) - oc = self.softmax(x) - B, C, T = oc.shape - tgrid = torch.arange(0, T, 1, device=device).unsqueeze(0) * 1 + torch.arange(0, batchlen, 1, device=device).unsqueeze(1) * batchstride - oc = oc.permute(0, 2, 1).reshape(-1, C) - oc = oc[:, [2, 0, 1]] - ot = tgrid.squeeze() - ot = ot.reshape(-1) - return oc, ot -model = Picker() + super().__init__(UNet, ckpt_path="ckpt/china.unet.pt", seqlen=3072, overlap=1536) + + +model = Picker() model.eval() -ckpt = torch.load("model_list/9_sc.pt", weights_only=False, map_location="cpu") -state = ckpt.state_dict() -model.load_state_dict(state) input_names = ["wave"] output_names = ["prob", "time"] -#x = torch.randn([10, 3, 6144, 1]) x = torch.randn([500000, 3]) -torch.onnx.export(model, x, -"pickers/9_sc.onnx", verbose=True, -dynamic_axes={"wave":{0:"batch"}, "prob":{0:"batch"}, "time":{0:"batch"}}, -input_names=input_names, output_names=output_names, opset_version=11) \ No newline at end of file +torch.onnx.export( + model, + x, + "pickers/unet.onnx", + verbose=True, + dynamic_axes={"wave": {0: "batch"}, "prob": {0: "batch"}, "time": {0: "batch"}}, + input_names=input_names, + output_names=output_names, + opset_version=11, +) diff --git a/makeonnx.unetpp.py b/makeonnx.unetpp.py index a01aebc..2749320 100755 --- a/makeonnx.unetpp.py +++ b/makeonnx.unetpp.py @@ -1,62 +1,25 @@ -import torch -import torch.nn as nn -from models.UNetPlusPlus import UNetpp -class Picker(UNetpp): - def __init__(self): - super().__init__() - self.n_stride = 1 - def forward(self, x): - device = x.device - with torch.no_grad(): - #print("数据维度", x.shape) - T, C = x.shape - seqlen = 6144 - batchstride = 6144 - 256 - batchlen = torch.ceil(torch.tensor(T / batchstride).to(device)) - idx = torch.arange(0, seqlen, 1, device=device).unsqueeze(0) + torch.arange(0, batchlen, 1, device=device).unsqueeze(1) * batchstride - idx = idx.clamp(min=0, max=T-1).long() - x = x.to(device) - wave = x[idx, :] - wave = wave.permute(0, 2, 1) - wave -= torch.mean(wave, dim=2, keepdim=True) - max, maxidx = torch.max(torch.abs(wave), dim=2, keepdim=True) - wave /= (max + 1e-6) - #x = wave.unsqueeze(3) - x0_0 = self.conv0_0(wave) - x1_0 = self.conv1_0(self.pool(x0_0)) - x0_1 = self.conv0_1(torch.cat([x0_0, self.up(x1_0)], 1)) +import torch +from models.UNetPlusPlus import UNetpp +from onnx_picker_base import OnnxSlidingWindowPicker - x2_0 = self.conv2_0(self.pool(x1_0)) - x1_1 = self.conv1_1(torch.cat([x1_0, self.up(x2_0)], 1)) - x0_2 = self.conv0_2(torch.cat([x0_0, x0_1, self.up(x1_1)], 1)) - x3_0 = self.conv3_0(self.pool(x2_0)) - x2_1 = self.conv2_1(torch.cat([x2_0, self.up(x3_0)], 1)) - x1_2 = self.conv1_2(torch.cat([x1_0, x1_1, self.up(x2_1)], 1)) - x0_3 = self.conv0_3(torch.cat([x0_0, x0_1, x0_2, self.up(x1_2)], 1)) +class Picker(OnnxSlidingWindowPicker): + def __init__(self): + super().__init__(UNetpp, ckpt_path="ckpt/china.unetpp.pt") - x4_0 = self.conv4_0(self.pool(x3_0)) - x3_1 = self.conv3_1(torch.cat([x3_0, self.up(x4_0)], 1)) - x2_2 = self.conv2_2(torch.cat([x2_0, x2_1, self.up(x3_1)], 1)) - x1_3 = self.conv1_3(torch.cat([x1_0, x1_1, x1_2, self.up(x2_2)], 1)) - x0_4 = self.conv0_4(torch.cat([x0_0, x0_1, x0_2, x0_3, self.up(x1_3)], 1)) - output = self.final(x0_4) - oc = self.sigmoid(output) - B, C, T = oc.shape - tgrid = torch.arange(0, T, 1, device=device).unsqueeze(0) * self.n_stride + torch.arange(0, batchlen, 1, device=device).unsqueeze(1) * batchstride - oc = oc.permute(0, 2, 1).reshape(-1, C) - ot = tgrid.squeeze() - ot = ot.reshape(-1) - return oc, ot -model = Picker() +model = Picker() model.eval() -model.load_state_dict(torch.load("ckpt/china.unetpp.pt", map_location="cpu")) input_names = ["wave"] output_names = ["prob", "time"] -#x = torch.randn([10, 3, 6144, 1]) x = torch.randn([500000, 3]) -torch.onnx.export(model, x, -"pickers/unetpp.onnx", verbose=True, -dynamic_axes={"wave":{0:"batch"}, "prob":{0:"batch"}, "time":{0:"batch"}}, -input_names=input_names, output_names=output_names, opset_version=11) \ No newline at end of file +torch.onnx.export( + model, + x, + "pickers/unetpp.onnx", + verbose=True, + dynamic_axes={"wave": {0: "batch"}, "prob": {0: "batch"}, "time": {0: "batch"}}, + input_names=input_names, + output_names=output_names, + opset_version=11, +) diff --git a/onnx_picker_base.py b/onnx_picker_base.py new file mode 100644 index 0000000..667e923 --- /dev/null +++ b/onnx_picker_base.py @@ -0,0 +1,74 @@ +import torch +import torch.nn as nn + + +class OnnxSlidingWindowPicker(nn.Module): + """ + Wrapper to standardize ONNX picker interfaces. + + Each picker stores its underlying network in ``self.model`` and handles + sliding-window preprocessing plus checkpoint loading. Checkpoints are + expected to prefix parameters with ``model.``; if the prefix is missing, + it is added automatically for backward compatibility. + """ + + def __init__( + self, + model_ctor, + ckpt_path=None, + *, + state_dict=None, + seqlen=6144, + overlap=256, + use_softmax=True, + ): + super().__init__() + self.model = model_ctor() + self.n_stride = 1 + self.seqlen = seqlen + self.batchstride = seqlen - overlap + self.use_softmax = use_softmax + + if state_dict is None: + if ckpt_path is None: + raise ValueError("Either ckpt_path or state_dict must be provided") + state_dict = torch.load(ckpt_path, map_location="cpu") + if not any(k.startswith("model.") for k in state_dict.keys()): + state_dict = {f"model.{k}": v for k, v in state_dict.items()} + self.load_state_dict(state_dict, strict=False) + + def forward(self, x): + device = x.device + with torch.no_grad(): + wave, batchlen = self.window_and_normalize(x, device) + + logits = self.model(wave) + if logits.dim() == 4: + logits = logits.squeeze(dim=3) + if self.use_softmax and logits.shape[1] > 1: + logits = logits.softmax(dim=1) + + B, C, T = logits.shape + tgrid = ( + torch.arange(0, T, 1, device=device).unsqueeze(0) * self.n_stride + + torch.arange(0, batchlen, 1, device=device).unsqueeze(1) + * self.batchstride + ) + oc = logits.permute(0, 2, 1).reshape(-1, C) + ot = tgrid.squeeze().reshape(-1) + return oc, ot + + def window_and_normalize(self, x, device): + T, _ = x.shape + batchlen = torch.ceil(torch.tensor(T / self.batchstride).to(device)) + idx = ( + torch.arange(0, self.seqlen, 1, device=device).unsqueeze(0) + + torch.arange(0, batchlen, 1, device=device).unsqueeze(1) * self.batchstride + ) + idx = idx.clamp(min=0, max=T - 1).long() + wave = x.to(device)[idx, :] + wave = wave.permute(0, 2, 1) + wave -= torch.mean(wave, dim=2, keepdim=True) + maxv, _ = torch.max(torch.abs(wave), dim=2, keepdim=True) + wave /= (maxv + 1e-6) + return wave, batchlen