Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+MP-SENet/best_ckpt/g_best_dns filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -23,7 +23,6 @@ build/
 *.swo
 
 # Project-specific
-MP-SENet/
 *.wav
 *.mp3
 *.flac

MP-SENet/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Yexin Lu
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

MP-SENet/README.md

@@ -0,0 +1,66 @@
+# Explicit Estimation of Magnitude and Phase Spectra in Parallel for High-Quality Speech Enhancement
+### Ye-Xin Lu, Yang Ai, Zhen-Hua Ling
+In our [paper](https://arxiv.org/abs/2305.13686), we proposed MP-SENet: a TF-domain monaural SE model with parallel magnitude and phase spectra denoising.<br>
+A [long-version](https://arxiv.org/abs/2308.08926) MP-SENet was extended to the speech denoising, dereverberation, and bandwidth extension tasks.<br>
+Audio samples can be found at the [demo website](http://yxlu-0102.github.io/MP-SENet).<br>
+We provide our implementation as open source in this repository.
+
+## ⚠️ Note
+There is a small bug in our code, but it does not affect the overall performance of the model. 
+If you intend to retrain the model, it’s **strongly recommended** to set `batch_first=True` in the `MultiHeadAttention` module inside [transformer.py](models/transformer.py), which can significantly reduce the memory usage of the model.
+
+## Pre-requisites
+1. Python >= 3.6.
+2. Clone this repository.
+3. Install python requirements. Please refer [requirements.txt](https://github.com/yxlu-0102/MP-SENet/blob/main/requirements.txt).
+4. Download and extract the [VoiceBank+DEMAND dataset](https://datashare.ed.ac.uk/handle/10283/1942). Resample all wav files to 16kHz, and move the clean and noisy wavs to `VoiceBank+DEMAND/wavs_clean` and `VoiceBank+DEMAND/wavs_noisy`, respectively. You can also directly download the downsampled 16kHz dataset [here](https://drive.google.com/drive/folders/19I_thf6F396y5gZxLTxYIojZXC0Ywm8l).
+
+## Training
+```
+CUDA_VISIBLE_DEVICES=0,1 python train.py --config config.json
+```
+Checkpoints and copy of the configuration file are saved in the `cp_mpsenet` directory by default.<br>
+You can change the path by adding `--checkpoint_path` option.
+
+## Inference
+```
+python inference.py --checkpoint_file [generator checkpoint file path]
+```
+You can also use the pretrained best checkpoint files we provide in the `best_ckpt` directory.
+<br>
+Generated wav files are saved in `generated_files` by default.
+You can change the path by adding `--output_dir` option.<br>
+Here is an example:
+```
+python inference.py --checkpoint_file best_ckpt/g_best_vb --output_dir generated_files/MP-SENet_VB
+```
+
+## Model Structure
+![model](Figures/model.png)
+
+## Comparison with other SE models
+![comparison](Figures/table.png)
+
+## Acknowledgements
+We referred to [HiFiGAN](https://github.com/jik876/hifi-gan), [NSPP](https://github.com/YangAi520/NSPP) 
+and [CMGAN](https://github.com/ruizhecao96/CMGAN) to implement this.
+
+## Citation
+```
+@inproceedings{lu2023mp,
+  title={{MP-SENet}: A Speech Enhancement Model with Parallel Denoising of Magnitude and Phase Spectra},
+  author={Lu, Ye-Xin and Ai, Yang and Ling, Zhen-Hua},
+  booktitle={Proc. Interspeech},
+  pages={3834--3838},
+  year={2023}
+}
+
+@article{lu2023explicit,
+  title={Explicit estimation of magnitude and phase spectra in parallel for high-quality speech enhancement},
+  author={Lu, Ye-Xin and Ai, Yang and Ling, Zhen-Hua},
+  journal={Neural Networks},
+  volume = {189},
+  pages = {107562},
+  year={2025}
+}
+```

MP-SENet/best_ckpt/config.json

@@ -0,0 +1,28 @@
+{
+    "num_gpus": 0,
+    "batch_size": 4,
+    "learning_rate": 0.0005,
+    "adam_b1": 0.8,
+    "adam_b2": 0.99,
+    "lr_decay": 0.99,
+    "seed": 1234,
+    
+    "dense_channel": 64,
+    "compress_factor": 0.3, 
+    "num_tsconformers": 4,
+    "beta": 2.0,
+
+    "sampling_rate": 16000,
+    "segment_size": 32000,
+    "n_fft": 400,
+    "hop_size": 100,
+    "win_size": 400,
+
+    "num_workers": 4,
+
+    "dist_config": {
+        "dist_backend": "nccl",
+        "dist_url": "tcp://localhost:54321",
+        "world_size": 1
+    }
+}

MP-SENet/best_ckpt/g_best_dns

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

MP-SENet/dataset.py

@@ -0,0 +1,95 @@
+import os
+import random
+import torch
+import torch.utils.data
+import librosa
+
+def mag_pha_stft(y, n_fft, hop_size, win_size, compress_factor=1.0, center=True):
+
+    hann_window = torch.hann_window(win_size).to(y.device)
+    stft_spec = torch.stft(y, n_fft, hop_length=hop_size, win_length=win_size, window=hann_window,
+                           center=center, pad_mode='reflect', normalized=False, return_complex=True)
+    stft_spec = torch.view_as_real(stft_spec)
+    mag = torch.sqrt(stft_spec.pow(2).sum(-1)+(1e-9))
+    pha = torch.atan2(stft_spec[:, :, :, 1]+(1e-10), stft_spec[:, :, :, 0]+(1e-5))
+    # Magnitude Compression
+    mag = torch.pow(mag, compress_factor)
+    com = torch.stack((mag*torch.cos(pha), mag*torch.sin(pha)), dim=-1)
+
+    return mag, pha, com
+
+
+def mag_pha_istft(mag, pha, n_fft, hop_size, win_size, compress_factor=1.0, center=True):
+    # Magnitude Decompression
+    mag = torch.pow(mag, (1.0/compress_factor))
+    com = torch.complex(mag*torch.cos(pha), mag*torch.sin(pha))
+    hann_window = torch.hann_window(win_size).to(com.device)
+    wav = torch.istft(com, n_fft, hop_length=hop_size, win_length=win_size, window=hann_window, center=center)
+
+    return wav
+
+
+def get_dataset_filelist(a):
+    with open(a.input_training_file, 'r', encoding='utf-8') as fi:
+        training_indexes = [x.split('|')[0] for x in fi.read().split('\n') if len(x) > 0]
+
+    with open(a.input_validation_file, 'r', encoding='utf-8') as fi:
+        validation_indexes = [x.split('|')[0] for x in fi.read().split('\n') if len(x) > 0]
+
+    return training_indexes, validation_indexes
+
+
+class Dataset(torch.utils.data.Dataset):
+    def __init__(self, training_indexes, clean_wavs_dir, noisy_wavs_dir, segment_size, 
+                sampling_rate, split=True, shuffle=True, n_cache_reuse=1, device=None):
+        self.audio_indexes = training_indexes
+        random.seed(1234)
+        if shuffle:
+            random.shuffle(self.audio_indexes)
+        self.clean_wavs_dir = clean_wavs_dir
+        self.noisy_wavs_dir = noisy_wavs_dir
+        self.segment_size = segment_size
+        self.sampling_rate = sampling_rate
+        self.split = split
+        self.cached_clean_wav = None
+        self.cached_noisy_wav = None
+        self.n_cache_reuse = n_cache_reuse
+        self._cache_ref_count = 0
+        self.device = device
+
+    def __getitem__(self, index):
+        filename = self.audio_indexes[index]
+        if self._cache_ref_count == 0:
+            clean_audio, _ = librosa.load(os.path.join(self.clean_wavs_dir, filename + '.wav'), sr=self.sampling_rate)
+            noisy_audio, _ = librosa.load(os.path.join(self.noisy_wavs_dir, filename + '.wav'), sr=self.sampling_rate)
+            length = min(len(clean_audio), len(noisy_audio))
+            clean_audio, noisy_audio = clean_audio[: length], noisy_audio[: length]
+            self.cached_clean_wav = clean_audio
+            self.cached_noisy_wav = noisy_audio
+            self._cache_ref_count = self.n_cache_reuse
+        else:
+            clean_audio = self.cached_clean_wav
+            noisy_audio = self.cached_noisy_wav
+            self._cache_ref_count -= 1
+        
+        clean_audio, noisy_audio = torch.FloatTensor(clean_audio), torch.FloatTensor(noisy_audio)
+        norm_factor = torch.sqrt(len(noisy_audio) / torch.sum(noisy_audio ** 2.0))
+        clean_audio = (clean_audio * norm_factor).unsqueeze(0)
+        noisy_audio = (noisy_audio * norm_factor).unsqueeze(0)
+
+        assert clean_audio.size(1) == noisy_audio.size(1)
+
+        if self.split:
+            if clean_audio.size(1) >= self.segment_size:
+                max_audio_start = clean_audio.size(1) - self.segment_size
+                audio_start = random.randint(0, max_audio_start)
+                clean_audio = clean_audio[:, audio_start: audio_start+self.segment_size]
+                noisy_audio = noisy_audio[:, audio_start: audio_start+self.segment_size]
+            else:
+                clean_audio = torch.nn.functional.pad(clean_audio, (0, self.segment_size - clean_audio.size(1)), 'constant')
+                noisy_audio = torch.nn.functional.pad(noisy_audio, (0, self.segment_size - noisy_audio.size(1)), 'constant')
+
+        return (clean_audio.squeeze(), noisy_audio.squeeze())
+
+    def __len__(self):
+        return len(self.audio_indexes)

MP-SENet/env.py

@@ -0,0 +1,14 @@
+import os
+import shutil
+
+class AttrDict(dict):
+    def __init__(self, *args, **kwargs):
+        super(AttrDict, self).__init__(*args, **kwargs)
+        self.__dict__ = self
+
+
+def build_env(config, config_name, path):
+    t_path = os.path.join(path, config_name)
+    if config != t_path:
+        os.makedirs(path, exist_ok=True)
+        shutil.copyfile(config, os.path.join(path, config_name))

MP-SENet/models/model.py

@@ -0,0 +1,195 @@
+import torch
+import torch.nn as nn
+import numpy as np
+from models.transformer import TransformerBlock
+from utils import LearnableSigmoid2d
+from pesq import pesq
+from joblib import Parallel, delayed
+
+class SPConvTranspose2d(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, r=1):
+        super(SPConvTranspose2d, self).__init__()
+        self.pad1 = nn.ConstantPad2d((1, 1, 0, 0), value=0.)
+        self.out_channels = out_channels
+        self.conv = nn.Conv2d(in_channels, out_channels * r, kernel_size=kernel_size, stride=(1, 1))
+        self.r = r
+
+    def forward(self, x):
+        x = self.pad1(x)
+        out = self.conv(x)
+        batch_size, nchannels, H, W = out.shape
+        out = out.view((batch_size, self.r, nchannels // self.r, H, W))
+        out = out.permute(0, 2, 3, 4, 1)
+        out = out.contiguous().view((batch_size, nchannels // self.r, H, -1))
+        return out
+    
+
+class DenseBlock(nn.Module):
+    def __init__(self, h, kernel_size=(2, 3), depth=4):
+        super(DenseBlock, self).__init__()
+        self.h = h
+        self.depth = depth
+        self.dense_block = nn.ModuleList([])
+        for i in range(depth):
+            dilation = 2 ** i
+            pad_length = dilation
+            dense_conv = nn.Sequential(
+                nn.ConstantPad2d((1, 1, pad_length, 0), value=0.),
+                nn.Conv2d(h.dense_channel*(i+1), h.dense_channel, kernel_size, dilation=(dilation, 1)),
+                nn.InstanceNorm2d(h.dense_channel, affine=True),
+                nn.PReLU(h.dense_channel)
+            )
+            self.dense_block.append(dense_conv)
+
+    def forward(self, x):
+        skip = x
+        for i in range(self.depth):
+            x = self.dense_block[i](skip)
+            skip = torch.cat([x, skip], dim=1)
+        return x
+
+
+class DenseEncoder(nn.Module):
+    def __init__(self, h, in_channel):
+        super(DenseEncoder, self).__init__()
+        self.h = h
+        self.dense_conv_1 = nn.Sequential(
+            nn.Conv2d(in_channel, h.dense_channel, (1, 1)),
+            nn.InstanceNorm2d(h.dense_channel, affine=True),
+            nn.PReLU(h.dense_channel))
+
+        self.dense_block = DenseBlock(h, depth=4)
+
+        self.dense_conv_2 = nn.Sequential(
+            nn.Conv2d(h.dense_channel, h.dense_channel, (1, 3), (1, 2), padding=(0, 1)),
+            nn.InstanceNorm2d(h.dense_channel, affine=True),
+            nn.PReLU(h.dense_channel))
+
+    def forward(self, x):
+        x = self.dense_conv_1(x)  # [b, 64, T, F]
+        x = self.dense_block(x)   # [b, 64, T, F]
+        x = self.dense_conv_2(x)  # [b, 64, T, F//2]
+        return x
+
+
+class MaskDecoder(nn.Module):
+    def __init__(self, h, out_channel=1):
+        super(MaskDecoder, self).__init__()
+        self.dense_block = DenseBlock(h, depth=4)
+        self.mask_conv = nn.Sequential(
+            SPConvTranspose2d(h.dense_channel, h.dense_channel, (1, 3), 2),
+            nn.InstanceNorm2d(h.dense_channel, affine=True),
+            nn.PReLU(h.dense_channel),
+            nn.Conv2d(h.dense_channel, out_channel, (1, 2))
+        )
+        self.lsigmoid = LearnableSigmoid2d(h.n_fft//2+1, beta=h.beta)
+
+    def forward(self, x):
+        x = self.dense_block(x)
+        x = self.mask_conv(x)
+        x = x.permute(0, 3, 2, 1).squeeze(-1) # [B, F, T]
+        x = self.lsigmoid(x)
+        return x
+
+
+class PhaseDecoder(nn.Module):
+    def __init__(self, h, out_channel=1):
+        super(PhaseDecoder, self).__init__()
+        self.dense_block = DenseBlock(h, depth=4)
+        self.phase_conv = nn.Sequential(
+            SPConvTranspose2d(h.dense_channel, h.dense_channel, (1, 3), 2),
+            nn.InstanceNorm2d(h.dense_channel, affine=True),
+            nn.PReLU(h.dense_channel)
+        )
+        self.phase_conv_r = nn.Conv2d(h.dense_channel, out_channel, (1, 2))
+        self.phase_conv_i = nn.Conv2d(h.dense_channel, out_channel, (1, 2))
+        
+    def forward(self, x):
+        x = self.dense_block(x)
+        x = self.phase_conv(x)
+        x_r = self.phase_conv_r(x)
+        x_i = self.phase_conv_i(x)
+        x = torch.atan2(x_i, x_r)
+        x = x.permute(0, 3, 2, 1).squeeze(-1) # [B, F, T]
+        return x
+
+
+class TSTransformerBlock(nn.Module):
+    def __init__(self, h):
+        super(TSTransformerBlock, self).__init__()
+        self.h = h
+        self.time_transformer = TransformerBlock(d_model=h.dense_channel, n_heads=4)
+        self.freq_transformer = TransformerBlock(d_model=h.dense_channel, n_heads=4)
+
+    def forward(self, x):
+        b, c, t, f = x.size()
+        x = x.permute(0, 3, 2, 1).contiguous().view(b*f, t, c)
+        x = self.time_transformer(x) + x
+        x = x.view(b, f, t, c).permute(0, 2, 1, 3).contiguous().view(b*t, f, c)
+        x = self.freq_transformer(x) + x
+        x = x.view(b, t, f, c).permute(0, 3, 1, 2)
+        return x
+
+
+class MPNet(nn.Module):
+    def __init__(self, h, num_tsblocks=4):
+        super(MPNet, self).__init__()
+        self.h = h
+        self.num_tscblocks = num_tsblocks
+        self.dense_encoder = DenseEncoder(h, in_channel=2)
+
+        self.TSTransformer = nn.ModuleList([])
+        for i in range(num_tsblocks):
+            self.TSTransformer.append(TSTransformerBlock(h))
+        
+        self.mask_decoder = MaskDecoder(h, out_channel=1)
+        self.phase_decoder = PhaseDecoder(h, out_channel=1)
+
+    def forward(self, noisy_amp, noisy_pha): # [B, F, T]
+
+        x = torch.stack((noisy_amp, noisy_pha), dim=-1).permute(0, 3, 2, 1) # [B, 2, T, F]
+        x = self.dense_encoder(x)
+
+        for i in range(self.num_tscblocks):
+            x = self.TSTransformer[i](x)
+        
+        denoised_amp = noisy_amp * self.mask_decoder(x)
+        denoised_pha = self.phase_decoder(x)
+        denoised_com = torch.stack((denoised_amp*torch.cos(denoised_pha),
+                                    denoised_amp*torch.sin(denoised_pha)), dim=-1)
+
+        return denoised_amp, denoised_pha, denoised_com
+
+
+def phase_losses(phase_r, phase_g):
+
+    ip_loss = torch.mean(anti_wrapping_function(phase_r - phase_g))
+    gd_loss = torch.mean(anti_wrapping_function(torch.diff(phase_r, dim=1) - torch.diff(phase_g, dim=1)))
+    iaf_loss = torch.mean(anti_wrapping_function(torch.diff(phase_r, dim=2) - torch.diff(phase_g, dim=2)))
+
+    return ip_loss, gd_loss, iaf_loss
+
+def anti_wrapping_function(x):
+
+    return torch.abs(x - torch.round(x / (2 * np.pi)) * 2 * np.pi)
+
+
+def pesq_score(utts_r, utts_g, h):
+
+    pesq_score = Parallel(n_jobs=30)(delayed(eval_pesq)(
+                            utts_r[i].squeeze().cpu().numpy(),
+                            utts_g[i].squeeze().cpu().numpy(), 
+                            h.sampling_rate)
+                          for i in range(len(utts_r)))
+    pesq_score = np.mean(pesq_score)
+
+    return pesq_score
+
+
+def eval_pesq(clean_utt, esti_utt, sr):
+    try:
+        pesq_score = pesq(sr, clean_utt, esti_utt)
+    except:
+        pesq_score = -1
+
+    return pesq_score

MP-SENet/models/transformer.py

@@ -0,0 +1,66 @@
+import torch
+import torch.nn.functional as F
+import torch.nn as nn
+from torch.nn import MultiheadAttention, GRU, Linear, LayerNorm, Dropout
+
+class FFN(nn.Module):
+    def __init__(self, d_model, bidirectional=True, dropout=0):
+        super(FFN, self).__init__()
+        self.gru = GRU(d_model, d_model*2, 1, bidirectional=bidirectional)
+        if bidirectional:
+            self.linear = Linear(d_model*2*2, d_model)
+        else:
+            self.linear = Linear(d_model*2, d_model)
+        self.dropout = Dropout(dropout)
+    
+    def forward(self, x):
+        self.gru.flatten_parameters()
+        x, _ = self.gru(x)
+        x = F.leaky_relu(x)
+        x = self.dropout(x)
+        x = self.linear(x)
+
+        return x
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, d_model, n_heads, bidirectional=True, dropout=0):
+        super(TransformerBlock, self).__init__()
+
+        self.norm1 = LayerNorm(d_model)
+        self.attention = MultiheadAttention(d_model, n_heads, dropout=dropout)
+        self.dropout1 = Dropout(dropout)
+        
+        self.norm2 = LayerNorm(d_model)
+        self.ffn = FFN(d_model, bidirectional=bidirectional)
+        self.dropout2 = Dropout(dropout)
+
+        self.norm3 = LayerNorm(d_model)
+
+    def forward(self, x, attn_mask=None, key_padding_mask=None):
+        xt = self.norm1(x)
+        xt, _ = self.attention(xt, xt, xt,
+                               attn_mask=attn_mask,
+                               key_padding_mask=key_padding_mask)
+        x = x + self.dropout1(xt)
+
+        xt = self.norm2(x)
+        xt = self.ffn(xt)
+        x = x + self.dropout2(xt)
+
+        x = self.norm3(x)
+
+        return x
+    
+def main():
+    x = torch.randn(4, 64, 401, 201)
+    b, c, t, f = x.size()
+    x = x.permute(0, 3, 2, 1).contiguous().view(b, f*t, c)
+    transformer = TransformerBlock(d_model=64, n_heads=4)
+    x = transformer(x)
+    x =  x.view(b, f, t, c).permute(0, 3, 2, 1)
+    print(x.size())
+
+
+if __name__ == '__main__':
+    main()

MP-SENet/utils.py

@@ -0,0 +1,99 @@
+import os
+import glob
+import torch
+import torch.nn as nn
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pylab as plt
+
+def plot_spectrogram(spectrogram):
+    fig, ax = plt.subplots(figsize=(4, 3))
+    im = ax.imshow(spectrogram, aspect="auto", origin="lower",
+                   interpolation='none')
+    plt.colorbar(im, ax=ax)
+    fig.canvas.draw()
+    plt.close()
+
+    return fig
+
+
+def get_padding(kernel_size, dilation=1):
+    return int((kernel_size*dilation - dilation)/2)
+
+
+def get_padding_2d(kernel_size, dilation=(1, 1)):
+    return (int((kernel_size[0]*dilation[0] - dilation[0])/2), int((kernel_size[1]*dilation[1] - dilation[1])/2))
+
+
+class LearnableSigmoid1d(nn.Module):
+    def __init__(self, in_features, beta=1):
+        super().__init__()
+        self.beta = beta
+        self.slope = nn.Parameter(torch.ones(in_features))
+        self.slope.requiresGrad = True
+
+    def forward(self, x):
+        return self.beta * torch.sigmoid(self.slope * x)
+
+
+class LearnableSigmoid2d(nn.Module):
+    def __init__(self, in_features, beta=1):
+        super().__init__()
+        self.beta = beta
+        self.slope = nn.Parameter(torch.ones(in_features, 1))
+        self.slope.requiresGrad = True
+
+    def forward(self, x):
+        return self.beta * torch.sigmoid(self.slope * x)
+
+
+class Sigmoid2d(nn.Module):
+    def __init__(self, in_features, beta=1):
+        super().__init__()
+        self.beta = beta
+        self.slope = torch.ones(in_features, 1)
+
+    def forward(self, x):
+        return self.beta * torch.sigmoid(self.slope * x)
+    
+
+class PLSigmoid(nn.Module):
+    def __init__(self, in_features):
+        super().__init__()
+        self.beta = nn.Parameter(torch.ones(in_features, 1) * 2.0)
+        self.slope = nn.Parameter(torch.ones(in_features, 1))
+        self.beta.requiresGrad = True
+        self.slope.requiresGrad = True
+
+    def forward(self, x):
+        return self.beta * torch.sigmoid(self.slope * x)
+    
+
+def load_checkpoint(filepath, device):
+    assert os.path.isfile(filepath)
+    print("Loading '{}'".format(filepath))
+    checkpoint_dict = torch.load(filepath, map_location=device)
+    print("Complete.")
+    return checkpoint_dict
+
+
+def save_checkpoint(filepath, obj):
+    print("Saving checkpoint to {}".format(filepath))
+    torch.save(obj, filepath)
+    print("Complete.")
+
+
+def scan_checkpoint(cp_dir, prefix):
+    pattern = os.path.join(cp_dir, prefix + '????????')
+    cp_list = glob.glob(pattern)
+    if len(cp_list) == 0:
+        return None
+    return sorted(cp_list)[-1]
+
+def main():
+    plsigmoid = PLSigmoid(201)
+    a = torch.randn(4, 201, 100)
+    print(plsigmoid(a))
+
+if __name__ == '__main__':
+    main()