Initial commit

app.py

@@ -0,0 +1,164 @@
+# Gradio app that takes seismic waveform as input and marks 2 phases on the waveform as output.
+
+import gradio as gr
+import numpy as np
+from phasehunter.model import Onset_picker, Updated_onset_picker
+from phasehunter.data_preparation import prepare_waveform
+import torch
+from scipy.stats.kde import gaussian_kde
+
+import obspy
+from obspy.clients.fdsn import Client
+from obspy.clients.fdsn.header import FDSNNoDataException, FDSNTimeoutException, FDSNInternalServerException
+from obspy.geodetics.base import locations2degrees
+from obspy.taup import TauPyModel
+from obspy.taup.helper_classes import SlownessModelError
+
+from obspy.clients.fdsn.header import URL_MAPPINGS
+
+def make_prediction(waveform):
+    waveform = np.load(waveform)
+    processed_input = prepare_waveform(waveform)
+    
+    # Make prediction
+    with torch.no_grad():
+        output = model(processed_input)
+
+    p_phase = output[:, 0]
+    s_phase = output[:, 1]
+
+    return processed_input, p_phase, s_phase
+
+def mark_phases(waveform):
+    processed_input, p_phase, s_phase = make_prediction(waveform)
+
+    # Create a plot of the waveform with the phases marked
+    if sum(processed_input[0][2] == 0): #if input is 1C
+        fig, ax = plt.subplots(nrows=2, figsize=(10, 2), sharex=True)
+
+        ax[0].plot(processed_input[0][0])
+        ax[0].set_ylabel('Norm. Ampl.')
+
+    else: #if input is 3C
+        fig, ax = plt.subplots(nrows=4, figsize=(10, 6), sharex=True)
+        ax[0].plot(processed_input[0][0])
+        ax[1].plot(processed_input[0][1])
+        ax[2].plot(processed_input[0][2])
+
+        ax[0].set_ylabel('Z')
+        ax[1].set_ylabel('N')
+        ax[2].set_ylabel('E')
+
+    p_phase_plot = p_phase*processed_input.shape[-1]
+    p_kde = gaussian_kde(p_phase_plot)
+    p_dist_space = np.linspace( min(p_phase_plot)-10, max(p_phase_plot)+10, 500 )
+    ax[-1].plot( p_dist_space, p_kde(p_dist_space), color='r')
+
+    s_phase_plot = s_phase*processed_input.shape[-1]
+    s_kde = gaussian_kde(s_phase_plot)
+    s_dist_space = np.linspace( min(s_phase_plot)-10, max(s_phase_plot)+10, 500 )
+    ax[-1].plot( s_dist_space, s_kde(s_dist_space), color='b')
+
+    for a in ax:
+        a.axvline(p_phase.mean()*processed_input.shape[-1], color='r', linestyle='--', label='P')
+        a.axvline(s_phase.mean()*processed_input.shape[-1], color='b', linestyle='--', label='S')
+
+    ax[-1].set_xlabel('Time, samples')
+    ax[-1].set_ylabel('Uncert.')
+    ax[-1].legend()
+
+    plt.subplots_adjust(hspace=0., wspace=0.)
+
+    # Convert the plot to an image and return it
+    fig.canvas.draw()
+    image = np.array(fig.canvas.renderer.buffer_rgba())
+    plt.close(fig)
+    return image
+
+def download_data(timestamp, eq_lat, eq_lon, client_name, radius_km):
+    client = Client(client_name)
+    window = radius_km / 111.2
+
+    assert eq_lat - window > -90 and eq_lat + window < 90, "Latitude out of bounds"
+    assert eq_lon - window > -180 and eq_lon + window < 180, "Longitude out of bounds"
+
+    # starttime = catalog['DateTime'].apply(lambda x: pd.Timestamp(x)).min()
+    # endtime = catalog['DateTime'].apply(lambda x: pd.Timestamp(x)).max()
+
+    return 0
+
+model = Onset_picker.load_from_checkpoint("./weights.ckpt",
+                                 picker=Updated_onset_picker(),
+                                    learning_rate=3e-4)
+model.eval()
+
+
+
+# # Create the Gradio interface
+# gr.Interface(mark_phases, inputs, outputs, title='PhaseHunter').launch()
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("# PhaseHunter")
+    with gr.Tab("Default example"):
+        # Define the input and output types for Gradio
+        inputs = gr.Dropdown(
+            ["data/sample/sample_0.npy", 
+            "data/sample/sample_1.npy", 
+            "data/sample/sample_2.npy"], 
+            label="Sample waveform", 
+            info="Select one of the samples",
+            value = "data/sample/sample_0.npy"
+        )
+
+        button = gr.Button("Predict phases")
+        outputs = gr.outputs.Image(label='Waveform with Phases Marked', type='numpy')
+    
+        button.click(mark_phases, inputs=inputs, outputs=outputs)
+        
+    with gr.Tab("Select earthquake from catalogue"):
+        gr.Markdown('TEST')
+        
+        client_inputs = gr.Dropdown(
+            choices = list(URL_MAPPINGS.keys()), 
+            label="FDSN Client", 
+            info="Select one of the available FDSN clients",
+            value = "IRIS",
+            interactive=True
+        )
+        with gr.Row(): 
+
+            timestamp_inputs = gr.Textbox(value='2019-07-04 17:33:49',
+                                placeholder='YYYY-MM-DD HH:MM:SS',
+                                label="Timestamp",
+                                info="Timestamp of the earthquake",
+                                max_lines=1,
+                                interactive=True)
+                               
+            eq_lat_inputs = gr.Number(value=35.766, 
+                            label="Latitude", 
+                            info="Latitude of the earthquake",
+                            interactive=True)
+            
+            eq_lo_inputs = gr.Number(value=117.605,
+                                label="Longitude",
+                                info="Longitude of the earthquake",
+                                interactive=True)
+            
+            radius_inputs = gr.Slider(minimum=1, 
+                                    maximum=150, 
+                                    value=50, label="Radius (km)", 
+                                    info="Select the radius around the earthquake to download data from",
+                                    interactive=True)
+            
+        button = gr.Button("Predict phases")
+
+    with gr.Tab("Predict on your own waveform"):
+        gr.Markdown("""
+        Please upload your waveform in .npy (numpy) format. 
+        Your waveform should be sampled at 100 sps and have 3 (Z, N, E) or 1 (Z) channels.
+        """)
+
+    button.click(mark_phases, inputs=inputs, outputs=outputs)
+
+demo.launch()

data/sample/sample_0.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bcacc216bea4273debc0244dce60def2413e4642100029fa0c0ce83416ba71c8
+size 144128

data/sample/sample_1.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:245f7bd592643f8573e8051aa593bdcdffa7e497de3e56aa76423cd96e44ae03
+size 113776

data/sample/sample_2.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c3eb64a6b01caba32464d950aafcbc4ed32f724f062c9c91b1c7f36671ce2b8e
+size 136344

data_preparation.py

@@ -0,0 +1,210 @@
+import torch
+import numpy as np
+
+from scipy import signal
+from scipy.signal import butter, lfilter, detrend
+
+# Make bandpass filter
+def butter_bandpass(lowcut, highcut, fs, order=5):
+    nyq = 0.5 * fs  # Nyquist frequency
+    low = lowcut / nyq  # Normalized frequency
+    high = highcut / nyq
+    b, a = butter(order, [low, high], btype="band")  # Bandpass filter
+    return b, a
+
+
+def butter_bandpass_filter(data, lowcut, highcut, fs, order=5):
+    b, a = butter_bandpass(lowcut, highcut, fs, order=order)
+    y = lfilter(b, a, data)
+    return y
+
+
+def rotate_waveform(waveform, angle):
+    fft_waveform = np.fft.fft(waveform)  # Compute the Fourier transform of the waveform
+    rotate_factor = np.exp(
+        1j * angle
+    )  # Create a complex exponential with the specified rotation angle
+    rotated_fft_waveform = (
+        fft_waveform * rotate_factor
+    )  # Multiply the Fourier transform by the rotation factor
+    rotated_waveform = np.fft.ifft(
+        rotated_fft_waveform
+    )  # Compute the inverse Fourier transform to get the rotated waveform in the time domain
+
+    return rotated_waveform
+
+
+def augment(sample):
+    # SET PARAMETERS:
+    crop_length = 6000
+    padding = 120
+    test = False
+
+    waveform = sample["waveform.npy"]
+    meta = sample["meta.json"]
+
+    if meta["split"] != "train":
+        test = True
+
+    target_sample_P = meta["trace_p_arrival_sample"]
+    target_sample_S = meta["trace_s_arrival_sample"]
+
+    if target_sample_P is None:
+        target_sample_P = 0
+    if target_sample_S is None:
+        target_sample_S = 0
+
+    # Randomly select a phase to start the crop
+    current_phases = [x for x in (target_sample_P, target_sample_S) if x > 0]
+    phase_selector = np.random.randint(0, len(current_phases))
+    first_phase = current_phases[phase_selector]
+
+    # Shuffle
+    if first_phase - (crop_length - padding) > padding:
+        start_indx = int(
+            first_phase
+            - torch.randint(low=padding, high=(crop_length - padding), size=(1,))
+        )
+        if test == True:
+            start_indx = int(first_phase - 2 * padding)
+
+    elif int(first_phase - padding) > 0:
+        start_indx = int(
+            first_phase
+            - torch.randint(low=0, high=(int(first_phase - padding)), size=(1,))
+        )
+        if test == True:
+            start_indx = int(first_phase - padding)
+
+    else:
+        start_indx = padding
+
+    end_indx = start_indx + crop_length
+
+    if (waveform.shape[-1] - end_indx) < 0:
+        start_indx += waveform.shape[-1] - end_indx
+        end_indx = start_indx + crop_length
+
+    # Update target
+    new_target_P = target_sample_P - start_indx
+    new_target_S = target_sample_S - start_indx
+
+    # Cut
+    waveform_cropped = waveform[:, start_indx:end_indx]
+
+    # Preprocess
+    waveform_cropped = detrend(waveform_cropped)
+    waveform_cropped = butter_bandpass_filter(
+        waveform_cropped, lowcut=0.2, highcut=40, fs=100, order=5
+    )
+    window = signal.windows.tukey(waveform_cropped[-1].shape[0], alpha=0.1)
+    waveform_cropped = waveform_cropped * window
+    waveform_cropped = detrend(waveform_cropped)
+
+    if np.isnan(waveform_cropped).any() == True:
+        waveform_cropped = np.zeros(shape=waveform_cropped.shape)
+
+        new_target_P = 0
+        new_target_S = 0
+
+    if np.sum(waveform_cropped) == 0:
+
+        new_target_P = 0
+        new_target_S = 0
+
+    # Normalize data
+    max_val = np.max(np.abs(waveform_cropped))
+    waveform_cropped_norm = waveform_cropped / max_val
+
+    # Added Z component only
+    if len(waveform_cropped_norm) < 3:
+        zeros = np.zeros((3, waveform_cropped_norm.shape[-1]))
+        zeros[0] = waveform_cropped_norm
+
+        waveform_cropped_norm = zeros
+
+    if test == False:
+        ##### Rotate waveform #####
+        probability = torch.randint(0, 2, size=(1,)).item()
+        angle = torch.FloatTensor(size=(1,)).uniform_(0.01, 359.9).item()
+        if probability == 1:
+            waveform_cropped_norm = rotate_waveform(waveform_cropped_norm, angle).real
+
+        #### Channel DropOUT #####
+        probability = torch.randint(0, 2, size=(1,)).item()
+        channel = torch.randint(1, 3, size=(1,)).item()
+        if probability == 1:
+            waveform_cropped_norm[channel, :] = 1e-6
+
+    # Normalize target
+    new_target_P = new_target_P / crop_length
+    new_target_S = new_target_S / crop_length
+
+    if (new_target_P <= 0) or (new_target_P >= 1) or (np.isnan(new_target_P)):
+        new_target_P = 0
+    if (new_target_S <= 0) or (new_target_S >= 1) or (np.isnan(new_target_S)):
+        new_target_S = 0
+
+    return waveform_cropped_norm, new_target_P, new_target_S
+
+
+def collation_fn(sample):
+    waveforms = np.stack([x[0] for x in sample])
+    targets_P = np.stack([x[1] for x in sample])
+    targets_S = np.stack([x[2] for x in sample])
+
+    return (
+        torch.tensor(waveforms, dtype=torch.float),
+        torch.tensor(targets_P, dtype=torch.float),
+        torch.tensor(targets_S, dtype=torch.float),
+    )
+
+
+def my_split_by_node(urls):
+    node_id, node_count = (
+        torch.distributed.get_rank(),
+        torch.distributed.get_world_size(),
+    )
+    return list(urls)[node_id::node_count]
+
+def prepare_waveform(waveform):
+    # SET PARAMETERS:
+    crop_length = 6000
+    padding = 120
+
+    assert waveform.shape[0] <= 3, "Waveform has more than 3 channels"
+
+    if waveform.shape[-1] < crop_length:
+        waveform = np.pad(
+            waveform,
+            ((0, 0), (0, crop_length - waveform.shape[-1])),
+            mode="constant",
+            constant_values=0,
+        )
+    if waveform.shape[-1] > crop_length:
+        waveform = waveform[:, :crop_length]
+
+    # Preprocess
+    waveform = detrend(waveform)
+    waveform = butter_bandpass_filter(
+        waveform, lowcut=0.2, highcut=40, fs=100, order=5
+    )
+    window = signal.windows.tukey(waveform[-1].shape[0], alpha=0.1)
+    waveform = waveform * window
+    waveform = detrend(waveform)
+
+    assert np.isnan(waveform).any() != True, "Nan in waveform"
+    assert np.sum(waveform) != 0, "Sum of waveform sample is zero"
+
+    # Normalize data
+    max_val = np.max(np.abs(waveform))
+    waveform = waveform / max_val
+
+    # Added Z component only
+    if len(waveform) < 3:
+        zeros = np.zeros((3, waveform.shape[-1]))
+        zeros[0] = waveform
+
+        waveform = zeros
+
+    return torch.tensor([waveform]*128, dtype=torch.float)

model.py

@@ -0,0 +1,313 @@
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torchmetrics import MeanAbsoluteError
+from torch.optim.lr_scheduler import ReduceLROnPlateau
+
+import lightning as pl
+
+class BlurPool1D(nn.Module):
+    def __init__(self, channels, pad_type="reflect", filt_size=3, stride=2, pad_off=0):
+        super(BlurPool1D, self).__init__()
+        self.filt_size = filt_size
+        self.pad_off = pad_off
+        self.pad_sizes = [
+            int(1.0 * (filt_size - 1) / 2),
+            int(np.ceil(1.0 * (filt_size - 1) / 2)),
+        ]
+        self.pad_sizes = [pad_size + pad_off for pad_size in self.pad_sizes]
+        self.stride = stride
+        self.off = int((self.stride - 1) / 2.0)
+        self.channels = channels
+
+        # print('Filter size [%i]' % filt_size)
+        if self.filt_size == 1:
+            a = np.array(
+                [
+                    1.0,
+                ]
+            )
+        elif self.filt_size == 2:
+            a = np.array([1.0, 1.0])
+        elif self.filt_size == 3:
+            a = np.array([1.0, 2.0, 1.0])
+        elif self.filt_size == 4:
+            a = np.array([1.0, 3.0, 3.0, 1.0])
+        elif self.filt_size == 5:
+            a = np.array([1.0, 4.0, 6.0, 4.0, 1.0])
+        elif self.filt_size == 6:
+            a = np.array([1.0, 5.0, 10.0, 10.0, 5.0, 1.0])
+        elif self.filt_size == 7:
+            a = np.array([1.0, 6.0, 15.0, 20.0, 15.0, 6.0, 1.0])
+
+        filt = torch.Tensor(a)
+        filt = filt / torch.sum(filt)
+        self.register_buffer("filt", filt[None, None, :].repeat((self.channels, 1, 1)))
+
+        self.pad = get_pad_layer_1d(pad_type)(self.pad_sizes)
+
+    def forward(self, inp):
+        if self.filt_size == 1:
+            if self.pad_off == 0:
+                return inp[:, :, :: self.stride]
+            else:
+                return self.pad(inp)[:, :, :: self.stride]
+        else:
+            return F.conv1d(
+                self.pad(inp), self.filt, stride=self.stride, groups=inp.shape[1]
+            )
+
+
+def get_pad_layer_1d(pad_type):
+    if pad_type in ["refl", "reflect"]:
+        PadLayer = nn.ReflectionPad1d
+    elif pad_type in ["repl", "replicate"]:
+        PadLayer = nn.ReplicationPad1d
+    elif pad_type == "zero":
+        PadLayer = nn.ZeroPad1d
+    else:
+        print("Pad type [%s] not recognized" % pad_type)
+    return PadLayer
+
+
+from masksembles import common
+
+
+class Masksembles1D(nn.Module):
+    def __init__(self, channels: int, n: int, scale: float):
+        super().__init__()
+
+        self.channels = channels
+        self.n = n
+        self.scale = scale
+
+        masks = common.generation_wrapper(channels, n, scale)
+        masks = torch.from_numpy(masks)
+
+        self.masks = torch.nn.Parameter(masks, requires_grad=False)
+
+    def forward(self, inputs):
+        batch = inputs.shape[0]
+        x = torch.split(inputs.unsqueeze(1), batch // self.n, dim=0)
+        x = torch.cat(x, dim=1).permute([1, 0, 2, 3])
+        x = x * self.masks.unsqueeze(1).unsqueeze(-1)
+        x = torch.cat(torch.split(x, 1, dim=0), dim=1)
+
+        return x.squeeze(0).type(inputs.dtype)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, in_planes, planes, stride=1, kernel_size=7, groups=1):
+        super(BasicBlock, self).__init__()
+        self.conv1 = nn.Conv1d(
+            in_planes,
+            planes,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding="same",
+            bias=False,
+        )
+        self.bn1 = nn.BatchNorm1d(planes)
+        self.conv2 = nn.Conv1d(
+            planes,
+            planes,
+            kernel_size=kernel_size,
+            stride=1,
+            padding="same",
+            bias=False,
+        )
+        self.bn2 = nn.BatchNorm1d(planes)
+
+        self.shortcut = nn.Sequential(
+            nn.Conv1d(
+                in_planes,
+                self.expansion * planes,
+                kernel_size=1,
+                stride=stride,
+                padding="same",
+                bias=False,
+            ),
+            nn.BatchNorm1d(self.expansion * planes),
+        )
+
+    def forward(self, x):
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = self.bn2(self.conv2(out))
+        out += self.shortcut(x)
+        out = F.relu(out)
+        return out
+
+
+class Updated_onset_picker(nn.Module):
+    def __init__(
+        self,
+    ):
+        super().__init__()
+
+        # self.activation = nn.ReLU()
+        # self.maxpool = nn.MaxPool1d(2)
+
+        self.n_masks = 128
+
+        self.block1 = nn.Sequential(
+            BasicBlock(3, 8, kernel_size=7, groups=1),
+            nn.GELU(),
+            BlurPool1D(8, filt_size=3, stride=2),
+            nn.GroupNorm(2, 8),
+        )
+
+        self.block2 = nn.Sequential(
+            BasicBlock(8, 16, kernel_size=7, groups=8),
+            nn.GELU(),
+            BlurPool1D(16, filt_size=3, stride=2),
+            nn.GroupNorm(2, 16),
+        )
+
+        self.block3 = nn.Sequential(
+            BasicBlock(16, 32, kernel_size=7, groups=16),
+            nn.GELU(),
+            BlurPool1D(32, filt_size=3, stride=2),
+            nn.GroupNorm(2, 32),
+        )
+
+        self.block4 = nn.Sequential(
+            BasicBlock(32, 64, kernel_size=7, groups=32),
+            nn.GELU(),
+            BlurPool1D(64, filt_size=3, stride=2),
+            nn.GroupNorm(2, 64),
+        )
+
+        self.block5 = nn.Sequential(
+            BasicBlock(64, 128, kernel_size=7, groups=64),
+            nn.GELU(),
+            BlurPool1D(128, filt_size=3, stride=2),
+            nn.GroupNorm(2, 128),
+        )
+
+        self.block6 = nn.Sequential(
+            Masksembles1D(128, self.n_masks, 2.0),
+            BasicBlock(128, 256, kernel_size=7, groups=128),
+            nn.GELU(),
+            BlurPool1D(256, filt_size=3, stride=2),
+            nn.GroupNorm(2, 256),
+        )
+
+        self.block7 = nn.Sequential(
+            Masksembles1D(256, self.n_masks, 2.0),
+            BasicBlock(256, 512, kernel_size=7, groups=256),
+            BlurPool1D(512, filt_size=3, stride=2),
+            nn.GELU(),
+            nn.GroupNorm(2, 512),
+        )
+
+        self.block8 = nn.Sequential(
+            Masksembles1D(512, self.n_masks, 2.0),
+            BasicBlock(512, 1024, kernel_size=7, groups=512),
+            BlurPool1D(1024, filt_size=3, stride=2),
+            nn.GELU(),
+            nn.GroupNorm(2, 1024),
+        )
+
+        self.block9 = nn.Sequential(
+            Masksembles1D(1024, self.n_masks, 2.0),
+            BasicBlock(1024, 128, kernel_size=7, groups=128),
+            # BlurPool1D(512, filt_size=3, stride=2),
+            # nn.GELU(),
+            # nn.GroupNorm(2,512),
+        )
+
+        self.out = nn.Sequential(nn.Linear(3072, 2), nn.Sigmoid())
+
+    def forward(self, x):
+        # Feature extraction
+
+        x = self.block1(x)
+        x = self.block2(x)
+
+        x = self.block3(x)
+        x = self.block4(x)
+
+        x = self.block5(x)
+        x = self.block6(x)
+
+        x = self.block7(x)
+        x = self.block8(x)
+
+        x = self.block9(x)
+
+        # Regressor
+        x = x.flatten(start_dim=1)
+        x = self.out(x)
+
+        return x
+
+class Onset_picker(pl.LightningModule):
+    def __init__(self, picker, learning_rate):
+        super().__init__()
+        self.picker = picker
+        self.learning_rate = learning_rate
+        self.save_hyperparameters(ignore=['picker'])
+        self.mae = MeanAbsoluteError()
+    
+    def compute_loss(self, y, pick, mae_name=False):
+        y_filt = y[y != 0]
+        pick_filt = pick[y != 0]
+        if len(y_filt) > 0:
+            loss = F.l1_loss(y_filt, pick_filt.flatten())
+            if mae_name != False:
+                mae_phase = self.mae(y_filt, pick_filt.flatten())*60
+                self.log(f'MAE/{mae_name}_val', mae_phase,  on_step=False, on_epoch=True, prog_bar=False, sync_dist=True)
+        else:
+            loss = 0
+        return loss
+            
+    def training_step(self, batch, batch_idx):
+        # training_step defines the train loop.
+        x, y_p, y_s = batch
+        # x, y_p, y_s, y_pg, y_sg, y_pn, y_sn = batch
+        
+        picks = self.picker(x)
+        
+        p_pick  = picks[:,0]
+        s_pick  = picks[:,1]
+        
+        p_loss = self.compute_loss(y_p, p_pick)
+        s_loss = self.compute_loss(y_s, s_pick)
+        
+        loss = (p_loss+s_loss)/2
+        
+        self.log('Loss/train', loss, on_step=True, on_epoch=False, prog_bar=True, sync_dist=True)
+        
+        return loss
+    
+    def validation_step(self, batch, batch_idx):
+        
+        x, y_p, y_s = batch
+        
+        picks = self.picker(x)
+        
+        p_pick  = picks[:,0]
+        s_pick  = picks[:,1]
+        
+        p_loss = self.compute_loss(y_p, p_pick, mae_name='P')
+        s_loss = self.compute_loss(y_s, s_pick, mae_name='S')
+        
+        loss = (p_loss+s_loss)/2
+            
+        self.log('Loss/val',  loss, on_step=False, on_epoch=True, prog_bar=False, sync_dist=True)
+        
+        return loss
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)
+        scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=10, cooldown=10, threshold=1e-3)
+        # scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, 3e-4,  epochs=300, steps_per_epoch=len(train_loader))
+        monitor = 'Loss/train'
+        return {"optimizer": optimizer,  "lr_scheduler": scheduler, 'monitor': monitor}
+    
+    def forward(self, x):
+        picks = self.picker(x)
+        return picks

training.py

@@ -0,0 +1,104 @@
+
+import torch
+
+from data_preparation import augment, collation_fn, my_split_by_node
+from model import Onset_picker, Updated_onset_picker
+
+import webdataset as wds
+
+from lightning.pytorch.callbacks import LearningRateMonitor, ModelCheckpoint
+from lightning.pytorch.loggers.tensorboard import TensorBoardLogger
+from lightning.pytorch.strategies import DDPStrategy
+from lightning import seed_everything
+import lightning as pl
+
+seed_everything(42, workers=False)
+torch.set_float32_matmul_precision('medium')
+
+batch_size = 256
+num_workers = 16 #int(os.cpu_count())
+n_iters_in_epoch = 5000
+
+train_dataset = (
+      wds.WebDataset("data/sample/shard-00{0000..0001}.tar", 
+                     # splitter=my_split_by_worker, 
+                     nodesplitter=my_split_by_node)
+      .decode()
+      .map(augment)
+      .shuffle(5000)
+      .batched(batchsize=batch_size,
+               collation_fn=collation_fn,
+               partial=False
+              )
+).with_epoch(n_iters_in_epoch//num_workers)
+
+
+val_dataset = (
+      wds.WebDataset("data/sample/shard-00{0000..0000}.tar", 
+                     # splitter=my_split_by_worker, 
+                     nodesplitter=my_split_by_node)
+      .decode()
+      .map(augment)
+      .repeat()
+      .batched(batchsize=batch_size,
+               collation_fn=collation_fn,
+               partial=False
+              )
+).with_epoch(100)
+
+
+train_loader = wds.WebLoader(train_dataset, 
+                             num_workers=num_workers, 
+                             shuffle=False,
+                             pin_memory=True, 
+                             batch_size=None)
+
+val_loader = wds.WebLoader(val_dataset, 
+                           num_workers=0,
+                           shuffle=False,
+                           pin_memory=True, 
+                           batch_size=None) 
+
+
+
+# model
+model = Onset_picker(picker=Updated_onset_picker(), 
+                     learning_rate=3e-4)
+# model = torch.compile(model, mode="reduce-overhead")
+
+logger = TensorBoardLogger("tensorboard_logdir", name="FAST")
+
+checkpoint_callback = ModelCheckpoint(save_top_k=1, monitor="Loss/val", filename="chkp-{epoch:02d}")
+lr_callback = LearningRateMonitor(logging_interval='epoch')
+# swa_callback = StochasticWeightAveraging(swa_lrs=0.05)
+    
+# # train model
+trainer = pl.Trainer(
+            precision='16-mixed',
+            
+            callbacks=[checkpoint_callback, lr_callback],
+    
+            devices='auto', 
+            accelerator='auto', 
+    
+            strategy=DDPStrategy(find_unused_parameters=False,
+                                 static_graph=True,
+                                 gradient_as_bucket_view=True),
+            benchmark=True,
+            
+            gradient_clip_val=0.5,
+            # ckpt_path='path/to/saved/checkpoints/chkp.ckpt',
+
+            # fast_dev_run=True,
+
+            logger=logger,
+            log_every_n_steps=50,
+            enable_progress_bar=True,
+    
+            max_epochs=300,
+        )
+
+trainer.fit(model=model, 
+            train_dataloaders=train_loader, 
+            val_dataloaders=val_loader,
+            )

weights.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:80255c65b749559f7c5c3f2bb993a25cc666d9a63a0d3050024679dd8064dcec
+size 200977197