Upload 1162 files

.gitattributes

@@ -33,3 +33,7 @@ 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
+s3prl_s3prl_main/s3prl/downstream/phone_linear/data/converted_aligned_phones.txt filter=lfs diff=lfs merge=lfs -text
+s3prl_s3prl_main/s3prl/downstream/voxceleb2_amsoftmax_segment_eval/cache_wav_paths/cache_test_segment.p filter=lfs diff=lfs merge=lfs -text
+s3prl_s3prl_main/s3prl/downstream/voxceleb2_amsoftmax_segment_eval/cache_wav_paths/cache_Voxceleb1.p filter=lfs diff=lfs merge=lfs -text
+s3prl_s3prl_main/s3prl/downstream/voxceleb2_amsoftmax_segment_eval/cache_wav_paths/cache_Voxceleb2.p filter=lfs diff=lfs merge=lfs -text

s3prl_s3prl_main/Dockerfile

@@ -0,0 +1,42 @@
+# We need this to use GPUs inside the container
+FROM nvidia/cuda:11.2.2-base
+# Using a multi-stage build simplifies the s3prl installation
+# TODO: Find a slimmer base image that also "just works"
+FROM tiangolo/uvicorn-gunicorn:python3.8
+
+RUN apt-get update --fix-missing && apt-get install -y wget \
+    libsndfile1 \
+    sox \
+    git \
+    git-lfs
+
+RUN python -m pip install --upgrade pip
+RUN python -m pip --no-cache-dir install fairseq@git+https://github.com//pytorch/fairseq.git@f2146bdc7abf293186de9449bfa2272775e39e1d#egg=fairseq
+RUN python -m pip --no-cache-dir install git+https://github.com/s3prl/s3prl.git#egg=s3prl
+
+COPY s3prl/ /app/s3prl
+COPY src/ /app/src
+
+# Setup filesystem
+RUN mkdir /app/data
+
+# Configure Git
+# TODO: Create a dedicated SUPERB account for the project?
+RUN git config --global user.email "[email protected]"
+RUN git config --global user.name "SUPERB Admin"
+
+# Default args for fine-tuning
+ENV upstream_model osanseviero/hubert_base
+ENV downstream_task asr
+ENV hub huggingface
+ENV hf_hub_org None
+ENV push_to_hf_hub True
+ENV override None
+
+WORKDIR /app/s3prl
+# Each task's config.yaml is used to set all the training parameters, but can be overridden with the `override` argument
+# The results of each training run are stored in /app/s3prl/result/downstream/{downstream_task}
+# and pushed to the Hugging Face Hub with name: 
+#   Default behaviour   - {hf_hub_username}/superb-s3prl-{upstream_model}-{downstream_task}-uuid
+#   With hf_hub_org set - {hf_hub_org}/superb-s3prl-{upstream_model}-{downstream_task}-uuid
+CMD python run_downstream.py -n ${downstream_task} -m train -u ${upstream_model} -d ${downstream_task} --hub ${hub} --hf_hub_org ${hf_hub_org} --push_to_hf_hub ${push_to_hf_hub} --override ${override}

s3prl_s3prl_main/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "{}"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2022 Andy T. Liu (Ting-Wei Liu) and Shu-wen (Leo) Yang
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

s3prl_s3prl_main/README.md

@@ -0,0 +1,357 @@
+<p align="center">
+    <img src="https://raw.githubusercontent.com/s3prl/s3prl/main/file/S3PRL-logo.png" width="900"/>
+    <br>
+    <br>
+    <a href="./LICENSE.txt"><img alt="Apache License 2.0" src="https://raw.githubusercontent.com/s3prl/s3prl/main/file/license.svg" /></a>
+    <a href="https://creativecommons.org/licenses/by-nc/4.0/"><img alt="CC_BY_NC License" src="https://img.shields.io/badge/License-CC%20BY--NC%204.0-lightgrey.svg" /></a>
+    <a href="https://github.com/s3prl/s3prl/actions/workflows/ci.yml"><img alt="CI" src="https://github.com/s3prl/s3prl/actions/workflows/ci.yml/badge.svg?branch=main&event=push"></a>
+    <a href="#development-pattern-for-contributors"><img alt="Codecov" src="https://img.shields.io/badge/contributions-welcome-brightgreen.svg"></a>
+    <a href="https://github.com/s3prl/s3prl/issues"><img alt="Bitbucket open issues" src="https://img.shields.io/github/issues/s3prl/s3prl"></a>
+</p>
+
+## Contact
+
+We prefer to have discussions directly on Github issue page, so that all the information is transparent to all the contributors and is auto-archived on the Github.
+If you wish to use email, please contact:
+
+- [Shu-wen (Leo) Yang](https://leo19941227.github.io/) ([email protected])
+- [Andy T. Liu](https://andi611.github.io/) ([email protected])
+
+Please refer to the [legacy citation](https://scholar.google.com/citations?view_op=view_citation&hl=en&user=R1mNI8QAAAAJ&citation_for_view=R1mNI8QAAAAJ:LkGwnXOMwfcC) of S3PRL and the timeline below, which justify our initiative on this project. This information is used to protect us from half-truths. We encourage to cite the individual papers most related to the function you are using to give fair credit to the developer of the function. You can find the names in the [Change Log](#change-log). Finally, we would like to thank our advisor, [Prof. Hung-yi Lee](https://speech.ee.ntu.edu.tw/~hylee/index.php), for his advice. The project would be impossible without his support.
+
+If you have any question (e.g., about who came up with / developed which ideas / functions or how the project started), feel free to engage in an open and responsible conversation on the GitHub issue page, and we'll be happy to help!
+
+## Contribution (pull request)
+
+**Guideline**
+
+- Starting in 2024, we will only accept new contributions in the form of new upstream models, so we can save bandwidth for developing new techniques (which will not be in S3PRL.)
+- S3PRL has transitioned into pure maintenance mode, ensuring the long-term maintenance of all existing functions.
+- Reporting bugs or the PR fixing the bugs is always welcome! Thanks!
+
+**Tutorials**
+
+- [General tutorial](https://s3prl.github.io/s3prl/contribute/general.html)
+- [Tutorial for adding new upstream models](https://s3prl.github.io/s3prl/contribute/upstream.html)
+
+## Environment compatibilities [![CI](https://github.com/s3prl/s3prl/actions/workflows/ci.yml/badge.svg?branch=main&event=push)](https://github.com/s3prl/s3prl/actions/workflows/ci.yml)
+
+We support the following environments. The test cases are ran with **[tox](./tox.ini)** locally and on **[github action](.github/workflows/ci.yml)**:
+
+| Env | versions |
+| --- | --- |
+| os  | `ubuntu-18.04`, `ubuntu-20.04` |
+| python | `3.7`, `3.8`, `3.9`, `3.10` |
+| pytorch | `1.8.1`, `1.9.1`, `1.10.2`, `1.11.0`, `1.12.1` , `1.13.1` , `2.0.1` , `2.1.0` |
+
+## Star History
+
+[![Star History Chart](https://api.star-history.com/svg?repos=s3prl/s3prl&type=Date)](https://star-history.com/#s3prl/s3prl&Date)
+
+## Change Log
+
+> We only list the major contributors here for conciseness. However, we are deeply grateful for all the contributions. Please see the [Contributors](https://github.com/s3prl/s3prl/graphs/contributors) page for the full list.
+
+* *Sep 2024*: Support MS-HuBERT (see [MS-HuBERT](https://arxiv.org/pdf/2406.05661))
+* *Dec 2023*: Support Multi-resolution HuBERT (MR-HuBERT, see [Multiresolution HuBERT](https://arxiv.org/pdf/2310.02720.pdf))
+* *Oct 2023*: Support ESPnet pre-trained upstream models (see [ESPnet HuBERT](https://arxiv.org/abs/2306.06672) and [WavLabLM](https://arxiv.org/abs/2309.15317))
+* *Sep 2022*: In [JSALT 2022](https://jsalt-2022-ssl.github.io/member), We upgrade the codebase to support testing, documentation and a new [S3PRL PyPI package](https://pypi.org/project/s3prl/) for easy installation and usage for upstream models. See our [online doc](https://s3prl.github.io/s3prl/) for more information. The package is now used by many [open-source projects](https://github.com/s3prl/s3prl/network/dependents), including [ESPNet](https://github.com/espnet/espnet/blob/master/espnet2/asr/frontend/s3prl.py). Contributors: [Shu-wen Yang](https://leo19941227.github.io/) ***(NTU)***, [Andy T. Liu](https://andi611.github.io/) ***(NTU)***, [Heng-Jui Chang](https://people.csail.mit.edu/hengjui/) ***(MIT)***, [Haibin Wu](https://hbwu-ntu.github.io/) ***(NTU)*** and [Xuankai Chang](https://www.xuankaic.com/) ***(CMU)***.
+* *Mar 2022*: Introduce [**SUPERB-SG**](https://arxiv.org/abs/2203.06849), see [Speech Translation](./s3prl/downstream/speech_translation) by [Hsiang-Sheng Tsai](https://github.com/bearhsiang) ***(NTU)***, [Out-of-domain ASR](./s3prl/downstream/ctc/) by [Heng-Jui Chang](https://people.csail.mit.edu/hengjui/) ***(NTU)***, [Voice Conversion](./s3prl/downstream/a2o-vc-vcc2020/) by [Wen-Chin Huang](https://unilight.github.io/) ***(Nagoya)***, [Speech Separation](./s3prl/downstream/separation_stft/) and [Speech Enhancement](./s3prl/downstream/enhancement_stft/) by [Zili Huang](https://scholar.google.com/citations?user=iQ-S0fQAAAAJ&hl=en) ***(JHU)*** for more info.
+* *Mar 2022*: Introduce [**SSL for SE/SS**](https://arxiv.org/abs/2203.07960) by [Zili Huang](https://scholar.google.com/citations?user=iQ-S0fQAAAAJ&hl=en) ***(JHU)***. See [SE1](https://github.com/s3prl/s3prl/tree/main/s3prl/downstream/enhancement_stft) and [SS1](https://github.com/s3prl/s3prl/tree/main/s3prl/downstream/separation_stft) folders for more details. Note that the improved performances can be achieved by the later introduced [SE2](https://github.com/s3prl/s3prl/tree/main/s3prl/downstream/enhancement_stft2) and [SS2](https://github.com/s3prl/s3prl/tree/main/s3prl/downstream/separation_stft2). However, for aligning with [SUPERB-SG](https://arxiv.org/abs/2203.06849) benchmarking, please use the version 1.
+* *Nov 2021*: Introduce [**S3PRL-VC**](https://arxiv.org/abs/2110.06280) by [Wen-Chin Huang](https://unilight.github.io/) ***(Nagoya)***, see [Any-to-one](https://github.com/s3prl/s3prl/tree/master/s3prl/downstream/a2o-vc-vcc2020) for more info. We highly recommend to consider the [newly released official repo of S3PRL-VC](https://github.com/unilight/s3prl-vc) which is developed and actively maintained by [Wen-Chin Huang](https://unilight.github.io/). The standalone repo contains much more recepies for the VC experiments. In S3PRL we only include the Any-to-one recipe for reproducing the SUPERB results.
+* *Oct 2021*: Support [**DistilHuBERT**](https://arxiv.org/abs/2110.01900) by [Heng-Jui Chang](https://people.csail.mit.edu/hengjui/) ***(NTU)***, see [docs](./s3prl/upstream/distiller/README.md) for more info.
+* *Sep 2021:* We host a *challenge* in [*AAAI workshop: The 2nd Self-supervised Learning for Audio and Speech Processing*](https://aaai-sas-2022.github.io/)! See [**SUPERB official site**](https://superbbenchmark.org/) for the challenge details and the [**SUPERB documentation**](./s3prl/downstream/docs/superb.md) in this toolkit!
+* *Aug 2021:* [Andy T. Liu](https://andi611.github.io/) ***(NTU)*** and [Shu-wen Yang](https://leo19941227.github.io/) ***(NTU)*** introduces the S3PRL toolkit in [MLSS 2021](https://ai.ntu.edu.tw/%e7%b7%9a%e4%b8%8a%e5%ad%b8%e7%bf%92-2/mlss-2021/), you can also **[watch it on Youtube](https://youtu.be/PkMFnS6cjAc)**!
+* *Aug 2021:* [**TERA**](https://ieeexplore.ieee.org/document/9478264) by [Andy T. Liu](https://andi611.github.io/) ***(NTU)*** is accepted to TASLP!
+* *July 2021:* We are now working on packaging s3prl and reorganizing the file structure in **v0.3**. Please consider using the stable **v0.2.0** for now. We will test and release **v0.3** before August.
+* *June 2021:* Support [**SUPERB:** **S**peech processing **U**niversal **PER**formance **B**enchmark](https://arxiv.org/abs/2105.01051), submitted to Interspeech 2021. Use the tag **superb-interspeech2021** or **v0.2.0**. Contributors: [Shu-wen Yang](https://leo19941227.github.io/) ***(NTU)***, [Pohan Chi](https://scholar.google.com/citations?user=SiyicoEAAAAJ&hl=zh-TW) ***(NTU)***, [Yist Lin](https://scholar.google.com/citations?user=0lrZq9MAAAAJ&hl=en) ***(NTU)***, [Yung-Sung Chuang](https://scholar.google.com/citations?user=3ar1DOwAAAAJ&hl=zh-TW) ***(NTU)***, [Jiatong Shi](https://scholar.google.com/citations?user=FEDNbgkAAAAJ&hl=en) ***(CMU)***, [Xuankai Chang](https://www.xuankaic.com/) ***(CMU)***, [Wei-Cheng Tseng](https://scholar.google.com.tw/citations?user=-d6aNP0AAAAJ&hl=zh-TW) ***(NTU)***, Tzu-Hsien Huang ***(NTU)*** and [Kushal Lakhotia](https://scholar.google.com/citations?user=w9W6zXUAAAAJ&hl=en) ***(Meta)***.
+* *June 2021:* Support extracting multiple hidden states for all the SSL pretrained models by [Shu-wen Yang](https://leo19941227.github.io/) ***(NTU)***.
+* *Jan 2021:* Readme updated with detailed instructions on how to use our latest version!
+* *Dec 2020:* We are migrating to a newer version for a more general, flexible, and scalable code. See the introduction below for more information! The legacy version can be accessed the tag **v0.1.0**.
+* *Oct 2020:* [Shu-wen Yang](https://leo19941227.github.io/) ***(NTU)*** and [Andy T. Liu](https://andi611.github.io/) ***(NTU)*** added varioius classic upstream models, including PASE+, APC, VQ-APC, NPC, wav2vec, vq-wav2vec ...etc.
+* *Oct 2019:* The birth of S3PRL! The repository was created for the [**Mockingjay**](https://arxiv.org/abs/1910.12638) development. [Andy T. Liu](https://andi611.github.io/) ***(NTU)***, [Shu-wen Yang](https://leo19941227.github.io/) ***(NTU)*** and [Pohan Chi](https://scholar.google.com/citations?user=SiyicoEAAAAJ&hl=zh-TW) ***(NTU)*** implemented the pre-training scripts and several simple downstream evaluation tasks. This work was the very start of the S3PRL project which established lots of foundamental modules and coding styles. Feel free to checkout to the old commits to explore our [legacy codebase](https://github.com/s3prl/s3prl/tree/6a53ee92bffeaa75fc2fb56071050bcf71e93785)!
+
+****
+
+## Introduction and Usages
+
+This is an open source toolkit called **s3prl**, which stands for **S**elf-**S**upervised **S**peech **P**re-training and **R**epresentation **L**earning.
+Self-supervised speech pre-trained models are called **upstream** in this toolkit, and are utilized in various **downstream** tasks.
+
+The toolkit has **three major usages**:
+
+### Pretrain
+
+- Pretrain upstream models, including Mockingjay, Audio ALBERT and TERA.
+- Document: [**pretrain/README.md**](./s3prl/pretrain/README.md)
+
+### Upstream
+
+- Easily load most of the existing upstream models with pretrained weights in a unified I/O interface.
+- Pretrained models are registered through **torch.hub**, which means you can use these models in your own project by one-line plug-and-play without depending on this toolkit's coding style.
+- Document: [**upstream/README.md**](./s3prl/upstream/README.md)
+
+### Downstream
+
+- Utilize upstream models in lots of downstream tasks
+- Benchmark upstream models with [**SUPERB Benchmark**](./s3prl/downstream/docs/superb.md)
+- Document: [**downstream/README.md**](./s3prl/downstream/README.md)
+
+---
+
+Here is a high-level illustration of how S3PRL might help you. We support to leverage numerous SSL representations on numerous speech processing tasks in our [GitHub codebase](https://github.com/s3prl/s3prl):
+
+![interface](file/S3PRL-interface.png)
+
+---
+
+We also modularize all the SSL models into a standalone [PyPi package](https://pypi.org/project/s3prl/) so that you can easily install it and use it without depending on our entire codebase. The following shows a simple example and you can find more details in our [documentation](https://s3prl.github.io/s3prl/).
+
+1. Install the S3PRL package:
+
+```sh
+pip install s3prl
+```
+
+2. Use it to extract representations for your own audio:
+
+```python
+import torch
+from s3prl.nn import S3PRLUpstream
+
+model = S3PRLUpstream("hubert")
+model.eval()
+
+with torch.no_grad():
+    wavs = torch.randn(2, 16000 * 2)
+    wavs_len = torch.LongTensor([16000 * 1, 16000 * 2])
+    all_hs, all_hs_len = model(wavs, wavs_len)
+
+for hs, hs_len in zip(all_hs, all_hs_len):
+    assert isinstance(hs, torch.FloatTensor)
+    assert isinstance(hs_len, torch.LongTensor)
+
+    batch_size, max_seq_len, hidden_size = hs.shape
+    assert hs_len.dim() == 1
+```
+
+---
+
+With this modularization, we have achieved close integration with the general speech processing toolkit [ESPNet](https://github.com/espnet/espnet), enabling the use of SSL models for a broader range of speech processing tasks and corpora to achieve state-of-the-art (SOTA) results (kudos to the [ESPNet Team](https://www.wavlab.org/open_source)):
+
+![integration](file/S3PRL-integration.png)
+
+You can start the journey of SSL with the following entry points:
+
+- S3PRL: [A simple SUPERB downstream task](https://github.com/s3prl/s3prl/blob/main/s3prl/downstream/docs/superb.md#pr-phoneme-recognition)
+- ESPNet: [Levearging S3PRL for ASR](https://github.com/espnet/espnet/tree/master/egs2/librispeech/asr1#self-supervised-learning-features-hubert_large_ll60k-conformer-utt_mvn-with-transformer-lm)
+
+---
+
+Feel free to use or modify our toolkit in your research. Here is a [list of papers using our toolkit](#used-by). Any question, bug report or improvement suggestion is welcome through [opening up a new issue](https://github.com/s3prl/s3prl/issues).
+
+If you find this toolkit helpful to your research, please do consider citing [our papers](#citation), thanks!
+
+## Installation
+
+1. **Python** >= 3.6
+2. Install **sox** on your OS
+3. Install s3prl: [Read doc](https://s3prl.github.io/s3prl/tutorial/installation.html#) or `pip install -e ".[all]"`
+4. (Optional) Some upstream models require special dependencies. If you encounter error with a specific upstream model, you can look into the `README.md` under each `upstream` folder. E.g., `upstream/pase/README.md`=
+
+## Reference Repositories
+
+* [Pytorch](https://github.com/pytorch/pytorch), Pytorch.
+* [Audio](https://github.com/pytorch/audio), Pytorch.
+* [Kaldi](https://github.com/kaldi-asr/kaldi), Kaldi-ASR.
+* [Transformers](https://github.com/huggingface/transformers), Hugging Face.
+* [PyTorch-Kaldi](https://github.com/mravanelli/pytorch-kaldi), Mirco Ravanelli.
+* [fairseq](https://github.com/pytorch/fairseq), Facebook AI Research.
+* [CPC](https://github.com/facebookresearch/CPC_audio), Facebook AI Research.
+* [APC](https://github.com/iamyuanchung/Autoregressive-Predictive-Coding), Yu-An Chung.
+* [VQ-APC](https://github.com/s3prl/VQ-APC), Yu-An Chung.
+* [NPC](https://github.com/Alexander-H-Liu/NPC), Alexander-H-Liu.
+* [End-to-end-ASR-Pytorch](https://github.com/Alexander-H-Liu/End-to-end-ASR-Pytorch), Alexander-H-Liu
+* [Mockingjay](https://github.com/andi611/Mockingjay-Speech-Representation), Andy T. Liu.
+* [ESPnet](https://github.com/espnet/espnet), Shinji Watanabe
+* [speech-representations](https://github.com/awslabs/speech-representations), aws lab
+* [PASE](https://github.com/santi-pdp/pase), Santiago Pascual and Mirco Ravanelli
+* [LibriMix](https://github.com/JorisCos/LibriMix), Joris Cosentino and Manuel Pariente
+
+## License
+
+The majority of S3PRL Toolkit is licensed under the Apache License version 2.0, however all the files authored by Facebook, Inc. (which have explicit copyright statement on the top) are licensed under CC-BY-NC.
+
+## Used by
+<details><summary>List of papers that used our toolkit (Feel free to add your own paper by making a pull request)</summary><p>
+
+### Self-Supervised Pretraining
+
++ [Mockingjay: Unsupervised Speech Representation Learning with Deep Bidirectional Transformer Encoders (Liu et al., 2020)](https://arxiv.org/abs/1910.12638)
+  ```
+  @article{mockingjay,
+     title={Mockingjay: Unsupervised Speech Representation Learning with Deep Bidirectional Transformer Encoders},
+     ISBN={9781509066315},
+     url={http://dx.doi.org/10.1109/ICASSP40776.2020.9054458},
+     DOI={10.1109/icassp40776.2020.9054458},
+     journal={ICASSP 2020 - 2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)},
+     publisher={IEEE},
+     author={Liu, Andy T. and Yang, Shu-wen and Chi, Po-Han and Hsu, Po-chun and Lee, Hung-yi},
+     year={2020},
+     month={May}
+  }
+  ```
++ [TERA: Self-Supervised Learning of Transformer Encoder Representation for Speech (Liu et al., 2020)](https://arxiv.org/abs/2007.06028)
+  ```
+  @misc{tera,
+      title={TERA: Self-Supervised Learning of Transformer Encoder Representation for Speech},
+      author={Andy T. Liu and Shang-Wen Li and Hung-yi Lee},
+      year={2020},
+      eprint={2007.06028},
+      archivePrefix={arXiv},
+      primaryClass={eess.AS}
+  }
+  ```
++ [Audio ALBERT: A Lite BERT for Self-supervised Learning of Audio Representation (Chi et al., 2020)](https://arxiv.org/abs/2005.08575)
+  ```
+  @inproceedings{audio_albert,
+      title={Audio ALBERT: A Lite BERT for Self-supervised Learning of Audio Representation},
+      author={Po-Han Chi and Pei-Hung Chung and Tsung-Han Wu and Chun-Cheng Hsieh and Shang-Wen Li and Hung-yi Lee},
+      year={2020},
+      booktitle={SLT 2020},
+  }
+  ```
+
+### Explanability
+
++ [Understanding Self-Attention of Self-Supervised Audio Transformers (Yang et al., 2020)](https://arxiv.org/abs/2006.03265)
+  ```
+  @inproceedings{understanding_sat,
+      author={Shu-wen Yang and Andy T. Liu and Hung-yi Lee},
+      title={{Understanding Self-Attention of Self-Supervised Audio Transformers}},
+      year=2020,
+      booktitle={Proc. Interspeech 2020},
+      pages={3785--3789},
+      doi={10.21437/Interspeech.2020-2231},
+      url={http://dx.doi.org/10.21437/Interspeech.2020-2231}
+  }
+  ```
+
+### Adversarial Attack
+
++ [Defense for Black-box Attacks on Anti-spoofing Models by Self-Supervised Learning (Wu et al., 2020)](https://arxiv.org/abs/2006.03214), code for computing LNSR: [utility/observe_lnsr.py](https://github.com/s3prl/s3prl/blob/master/utility/observe_lnsr.py)
+  ```
+  @inproceedings{mockingjay_defense,
+      author={Haibin Wu and Andy T. Liu and Hung-yi Lee},
+      title={{Defense for Black-Box Attacks on Anti-Spoofing Models by Self-Supervised Learning}},
+      year=2020,
+      booktitle={Proc. Interspeech 2020},
+      pages={3780--3784},
+      doi={10.21437/Interspeech.2020-2026},
+      url={http://dx.doi.org/10.21437/Interspeech.2020-2026}
+  }
+  ```
+
++ [Adversarial Defense for Automatic Speaker Verification by Cascaded Self-Supervised Learning Models (Wu et al., 2021)](https://arxiv.org/abs/2102.07047)
+  ```
+  @misc{asv_ssl,
+      title={Adversarial defense for automatic speaker verification by cascaded self-supervised learning models},
+      author={Haibin Wu and Xu Li and Andy T. Liu and Zhiyong Wu and Helen Meng and Hung-yi Lee},
+      year={2021},
+      eprint={2102.07047},
+      archivePrefix={arXiv},
+      primaryClass={eess.AS}
+  ```
+
+### Voice Conversion
+
++ [S2VC: A Framework for Any-to-Any Voice Conversion with Self-Supervised Pretrained Representations (Lin et al., 2021)](https://arxiv.org/abs/2104.02901)
+  ```
+  @misc{s2vc,
+        title={S2VC: A Framework for Any-to-Any Voice Conversion with Self-Supervised Pretrained Representations},
+        author={Jheng-hao Lin and Yist Y. Lin and Chung-Ming Chien and Hung-yi Lee},
+        year={2021},
+        eprint={2104.02901},
+        archivePrefix={arXiv},
+        primaryClass={eess.AS}
+  }
+  ```
+
+### Benchmark and Evaluation
+
++ [SUPERB: Speech processing Universal PERformance Benchmark (Yang et al., 2021)](https://arxiv.org/abs/2105.01051)
+  ```
+  @misc{superb,
+        title={SUPERB: Speech processing Universal PERformance Benchmark},
+        author={Shu-wen Yang and Po-Han Chi and Yung-Sung Chuang and Cheng-I Jeff Lai and Kushal Lakhotia and Yist Y. Lin and Andy T. Liu and Jiatong Shi and Xuankai Chang and Guan-Ting Lin and Tzu-Hsien Huang and Wei-Cheng Tseng and Ko-tik Lee and Da-Rong Liu and Zili Huang and Shuyan Dong and Shang-Wen Li and Shinji Watanabe and Abdelrahman Mohamed and Hung-yi Lee},
+        year={2021},
+        eprint={2105.01051},
+        archivePrefix={arXiv},
+        primaryClass={cs.CL}
+  }
+  ```
+
++ [Utilizing Self-supervised Representations for MOS Prediction (Tseng et al., 2021)](https://arxiv.org/abs/2104.03017)
+  ```
+  @misc{ssr_mos,
+      title={Utilizing Self-supervised Representations for MOS Prediction},
+      author={Wei-Cheng Tseng and Chien-yu Huang and Wei-Tsung Kao and Yist Y. Lin and Hung-yi Lee},
+      year={2021},
+      eprint={2104.03017},
+      archivePrefix={arXiv},
+      primaryClass={eess.AS}
+  }
+  ```
+}
+
+</p></details>
+
+## Citation
+
+If you find this toolkit useful, please consider citing following papers.
+
+- If you use our pre-training scripts, or the downstream tasks considered in *TERA* and *Mockingjay*, please consider citing the following:
+```
+@misc{tera,
+  title={TERA: Self-Supervised Learning of Transformer Encoder Representation for Speech},
+  author={Andy T. Liu and Shang-Wen Li and Hung-yi Lee},
+  year={2020},
+  eprint={2007.06028},
+  archivePrefix={arXiv},
+  primaryClass={eess.AS}
+}
+```
+```
+@article{mockingjay,
+   title={Mockingjay: Unsupervised Speech Representation Learning with Deep Bidirectional Transformer Encoders},
+   ISBN={9781509066315},
+   url={http://dx.doi.org/10.1109/ICASSP40776.2020.9054458},
+   DOI={10.1109/icassp40776.2020.9054458},
+   journal={ICASSP 2020 - 2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)},
+   publisher={IEEE},
+   author={Liu, Andy T. and Yang, Shu-wen and Chi, Po-Han and Hsu, Po-chun and Lee, Hung-yi},
+   year={2020},
+   month={May}
+}
+```
+
+- If you use our organized upstream interface and features, or the *SUPERB* downstream benchmark, please consider citing the following:
+```
+@article{yang2024large,
+  title={A Large-Scale Evaluation of Speech Foundation Models},
+  author={Yang, Shu-wen and Chang, Heng-Jui and Huang, Zili and Liu, Andy T and Lai, Cheng-I and Wu, Haibin and Shi, Jiatong and Chang, Xuankai and Tsai, Hsiang-Sheng and Huang, Wen-Chin and others},
+  journal={IEEE/ACM Transactions on Audio, Speech, and Language Processing},
+  year={2024},
+  publisher={IEEE}
+}
+```
+```
+@inproceedings{yang21c_interspeech,
+  author={Shu-wen Yang and Po-Han Chi and Yung-Sung Chuang and Cheng-I Jeff Lai and Kushal Lakhotia and Yist Y. Lin and Andy T. Liu and Jiatong Shi and Xuankai Chang and Guan-Ting Lin and Tzu-Hsien Huang and Wei-Cheng Tseng and Ko-tik Lee and Da-Rong Liu and Zili Huang and Shuyan Dong and Shang-Wen Li and Shinji Watanabe and Abdelrahman Mohamed and Hung-yi Lee},
+  title={{SUPERB: Speech Processing Universal PERformance Benchmark}},
+  year=2021,
+  booktitle={Proc. Interspeech 2021},
+  pages={1194--1198},
+  doi={10.21437/Interspeech.2021-1775}
+}
+```

s3prl_s3prl_main/ci/format.py

@@ -0,0 +1,84 @@
+#!/usr/bin/env python3
+
+import argparse
+from pathlib import Path
+from subprocess import CalledProcessError, check_output
+
+
+def load_valid_paths():
+    with open("./valid_paths.txt", "r") as fp:
+        paths = [line.strip() for line in fp if line.strip() != ""]
+        return paths
+
+
+def get_third_party():
+    txt_files = list(Path("./requirements").rglob("*.txt"))
+    package_list = []
+    for file in txt_files:
+        with open(file, "r") as fp:
+            for line in fp:
+                line = line.strip()
+                if line == "":
+                    continue
+                package_list.append(line.split(" ")[0])
+    return package_list
+
+
+def run_command(command: str):
+    try:
+        check_output(command.split(" "))
+    except CalledProcessError as e:
+        print(e.output.decode("utf-8"))
+        raise
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "files",
+        type=str,
+        nargs="*",
+        default=[],
+        help="If no file is given, use the files under ./valid_paths.txt",
+    )
+    parser.add_argument("--check", action="store_true", help="Only checks the files")
+    args = parser.parse_args()
+
+    if len(args.files) == 0:
+        args.files = load_valid_paths()
+
+    print(f"Formatting files: {args.files}")
+    args.files = " ".join(args.files)
+
+    print("Run flake8")
+    # stop the build if there are Python syntax errors or undefined names
+    run_command(
+        f"flake8 {args.files} --count --select=E9,F63,F7,F82 --show-source --statistics"
+    )
+    # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
+    run_command(
+        f"flake8 {args.files} --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics"
+    )
+
+    print("Run black")
+    if args.check:
+        run_command(f"black --check {args.files}")
+    else:
+        run_command(f"black {args.files}")
+
+    print("Run isort")
+    third_party = get_third_party()
+    third_party = ",".join(third_party)
+    if args.check:
+        run_command(
+            f"isort --profile black --thirdparty {third_party} --check {args.files}"
+        )
+    else:
+        run_command(f"isort --profile black --thirdparty {third_party} {args.files}")
+
+    if args.check:
+        print("Successfully passed the format check!")
+
+
+if __name__ == "__main__":
+    main()

s3prl_s3prl_main/docs/Makefile

@@ -0,0 +1,20 @@
+# Minimal makefile for Sphinx documentation
+#
+
+# You can set these variables from the command line, and also
+# from the environment for the first two.
+SPHINXOPTS    ?=
+SPHINXBUILD   ?= sphinx-build
+SOURCEDIR     = source
+BUILDDIR      = build
+
+# Put it first so that "make" without argument is like "make help".
+help:
+	@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
+
+.PHONY: help Makefile
+
+# Catch-all target: route all unknown targets to Sphinx using the new
+# "make mode" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).
+%: Makefile
+	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)

s3prl_s3prl_main/docs/README.md

@@ -0,0 +1,27 @@
+# Documentation
+
+To auto-generate documents for S3PRL, please follow the following steps:
+
+1. Activate an python env for the doc-generating tool to import all the modules to realize auto-documentation. There are various ways to achieve this. You can also follow:
+
+```sh
+conda create -y -n doc python=3.8
+conda activate doc
+
+cd $S3PRL_ROOT
+pip install ".[dev]"
+```
+
+2. Auto-generate HTML files for all the packages, modules and their submodules listed in `$S3PRL_ROOT/valid_paths.txt`. The HTML files will appear in `$S3PRL_ROOT/docs/build/html`
+
+```sh
+cd $S3PRL_ROOT/docs_src
+./rebuild_docs.sh
+```
+
+3. Launch the simple webserver to see the documentation.
+
+```sh
+cd $S3PRL_ROOT/docs/build/html
+python3 -m http.server
+```

s3prl_s3prl_main/docs/from_scratch_tutorial.md

@@ -0,0 +1,146 @@
+# REDAME
+## install sphinx
+```shell
+$ pip install -U sphinx
+```
+
+## set up
+at s3prl root directory (s3prl-private/)
+```shell
+$ mkdir docs
+$ cd docs
+$ sphinx-quickstart
+$ cd ..
+$ sphinx-apidoc -d 3 --separate --implicit-namespace -o docs ./s3prl s3prl/downstream s3prl/interface s3prl/preprocess s3prl/pretrain s3prl/problem s3prl/sampler s3prl/submit s3prl/superb s3prl/upstream s3prl/utility s3prl/wrapper s3prl/__init__.py s3prl/hub.py s3prl/optimizers.py s3prl/run_downstream.py s3prl/run_pretrain.py s3prl/run_while.sh s3prl/schedulers.py
+```
+
+## install dependencies
+```shell
+$ cd s3prl-private/docs
+$ echo "furo
+torch
+numpy
+pandas
+tqdm
+pytorch_lightning
+matplotlib
+ipdb>=0.13.9
+PyYAML
+transformers
+torchaudio
+gdown
+sklearn
+joblib
+tensorboardX
+librosa
+scipy
+lxml
+h5py
+dtw
+catalyst
+sox
+six
+easydict
+Resemblyzer
+sentencepiece
+pysoundfile
+asteroid
+sacrebleu
+speechbrain
+omegaconf
+editdistance" > requirement.txt
+$ pip install -r requirement.txt
+```
+
+add custom.js at s3prl-private/docs/_static/js/
+(just paste the following lines)
+```javascript=
+/*
+change the default sphinx.ext.linkcode's [source] to [Github]
+*/
+document.querySelectorAll(".reference.external .viewcode-link .pre").forEach(item => {
+    item.innerHTML = "[Github]"
+    item.style.marginRight = "3px"
+})
+
+```
+
+modify s3prl-private/docs/index.rst
+```diff
+# remove these lines
+- .. toctree::
+-    :maxdepth: 2
+-    :caption: Contents:
+
+# replaced with this line
++ .. include:: s3prl.rst
+```
+
+modify s3prl-private/docs/conf.py
+```python
+# add these lines at top
+import inspect
+import os
+import sys
+for x in os.walk('..'):
+	sys.path.insert(0, x[0])
+    
+# add extensions
+extensions = [
+    'sphinx.ext.autodoc',
+    'sphinx.ext.napoleon',
+    'sphinx.ext.viewcode',
+    'sphinx.ext.linkcode'
+]
+
+html_js_files = [
+    'js/custom.js',
+]
+
+def linkcode_resolve(domain, info):
+    def find_source():
+        obj = sys.modules[info['module']]
+        if info['fullname'] == 'InitConfig.args':	return None	
+        if info['fullname'] == 'InitConfig.kwargs':	return None	
+        for part in info['fullname'].split('.'):
+            obj = getattr(obj, part) 
+
+        if isinstance(obj, property):	return None	
+
+        fn = inspect.getsourcefile(obj)	
+        fn = os.path.relpath(fn, start=os.path.dirname(os.path.abspath(__file__))[:-4])
+
+        source, lineno = inspect.getsourcelines(obj)
+        return fn, lineno, lineno + len(source) - 1
+
+    if domain != 'py' or not info['module']:	return None
+
+    tag = 'master' if 'dev' in release else ('v' + release)		# s3prl github version
+
+    try:
+        filename = '%s#L%d-L%d' % find_source()			# specify file page with line number
+    except Exception:
+        filename = info['module'].replace('.', '/') + '.py'		# cannot find corresponding codeblock, use the file page instead
+
+    return "https://github.com/s3prl/s3prl-private/blob/%s/%s" % (tag, filename)
+```
+
+to use the furo theme, add this line in s3prl-private/docs/conf.py (replace the original alabaster theme)
+```python
+html_theme = "furo"
+```
+
+## generate html files
+at s3prl-private/docs/
+```shell
+$ make html
+```
+
+the html files will be generated at **s3prl-private/docs/_build/html/**
+click on **index.html** to view the doc page on your browser
+
+if you want to see how your modified codes looks like, simply do
+```shell
+$ make clean html        # this remove the old html files
+$ make html              # generate new html files
+```

s3prl_s3prl_main/docs/make.bat

@@ -0,0 +1,35 @@
+@ECHO OFF
+
+pushd %~dp0
+
+REM Command file for Sphinx documentation
+
+if "%SPHINXBUILD%" == "" (
+	set SPHINXBUILD=sphinx-build
+)
+set SOURCEDIR=source
+set BUILDDIR=build
+
+%SPHINXBUILD% >NUL 2>NUL
+if errorlevel 9009 (
+	echo.
+	echo.The 'sphinx-build' command was not found. Make sure you have Sphinx
+	echo.installed, then set the SPHINXBUILD environment variable to point
+	echo.to the full path of the 'sphinx-build' executable. Alternatively you
+	echo.may add the Sphinx directory to PATH.
+	echo.
+	echo.If you don't have Sphinx installed, grab it from
+	echo.https://www.sphinx-doc.org/
+	exit /b 1
+)
+
+if "%1" == "" goto help
+
+%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
+goto end
+
+:help
+%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
+
+:end
+popd

s3prl_s3prl_main/docs/rebuild_docs.sh

@@ -0,0 +1,9 @@
+#!/bin/bash
+
+rm -rf ./source/s3prl.*
+rm -rf ./source/_autosummary
+
+make clean html
+
+touch build/html/.nojekyll
+

s3prl_s3prl_main/docs/source/_static/css/custom.css

@@ -0,0 +1,3 @@
+.toctree-wrapper.compound .caption-text {
+    display: none;
+}

s3prl_s3prl_main/docs/source/_static/js/custom.js

@@ -0,0 +1,7 @@
+/*
+change the default sphinx.ext.linkcode's [source] to [Github]
+*/
+document.querySelectorAll(".reference.external .viewcode-link .pre").forEach(item => {
+    item.innerHTML = "[Github]"
+    item.style.marginRight = "3px"
+})

s3prl_s3prl_main/docs/source/_templates/custom-module-template.rst

@@ -0,0 +1,81 @@
+{{ name | escape | underline}}
+
+**({{ fullname }})**
+
+.. automodule:: {{ fullname }}
+
+{% block modules %}
+{% if modules %}
+
+.. autosummary::
+   :toctree:
+   :template: custom-module-template.rst
+   :recursive:
+
+{% for item in modules %}
+   {{ item }}
+{%- endfor %}
+{% endif %}
+{% endblock %}
+
+
+{% block attributes %}
+{% if attributes %}
+{% for item in attributes %}
+
+{{ item }}
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoattribute:: {{ item }}
+
+{%- endfor %}
+{% endif %}
+{% endblock %}
+
+{% block classes %}
+{% if classes %}
+{% for item in classes %}
+
+{{ item }}
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: {{ item }}
+   :member-order: bysource
+   :members:
+   :undoc-members:
+   :inherited-members: torch.nn.Module,nn.Module,Module
+   :show-inheritance:
+
+{%- endfor %}
+{% endif %}
+{% endblock %}
+
+{% block functions %}
+{% if functions %}
+{% for item in functions %}
+
+{{ item }}
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: {{ item }}
+
+{%- endfor %}
+{% endif %}
+{% endblock %}
+
+{% block exceptions %}
+{% if exceptions %}
+{% for item in exceptions %}
+
+{{ item }}
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoexception:: {{ item }}
+   :members:
+   :undoc-members:
+   :show-inheritance:
+   :member-order: bysource
+
+{%- endfor %}
+{% endif %}
+{% endblock %}

s3prl_s3prl_main/docs/source/conf.py

@@ -0,0 +1,120 @@
+# Configuration file for the Sphinx documentation builder.
+#
+# This file only contains a selection of the most common options. For a full
+# list see the documentation:
+# https://www.sphinx-doc.org/en/master/usage/configuration.html
+
+# -- Path setup --------------------------------------------------------------
+
+# If extensions (or modules to document with autodoc) are in another directory,
+# add these directories to sys.path here. If the directory is relative to the
+# documentation root, use os.path.abspath to make it absolute, like shown here.
+#
+import inspect
+import os
+import sys
+from pathlib import Path
+
+# -- Project information -----------------------------------------------------
+
+project = "S3PRL"
+copyright = "2022, S3PRL Team"
+author = "S3PRL Team"
+
+# The full version, including alpha/beta/rc tags
+with (Path(__file__).parent.parent.parent / "s3prl" / "version.txt").open() as f:
+    release = f.read()
+
+
+def linkcode_resolve(domain, info):
+    def find_source():
+        obj = sys.modules[info["module"]]
+        for part in info["fullname"].split("."):
+            obj = getattr(obj, part)
+
+        if isinstance(obj, property):
+            return None
+
+        file_parts = Path(inspect.getsourcefile(obj)).parts
+        reversed_parts = []
+        for part in reversed(file_parts):
+            if part == "s3prl":
+                reversed_parts.append(part)
+                break
+            else:
+                reversed_parts.append(part)
+        fn = "/".join(reversed(reversed_parts))
+
+        source, lineno = inspect.getsourcelines(obj)
+        return fn, lineno, lineno + len(source) - 1
+
+    if domain != "py" or not info["module"]:
+        return None
+
+    tag = "master" if "dev" in release else ("v" + release)  # s3prl github version
+
+    try:
+        filename = "%s#L%d-L%d" % find_source()  # specify file page with line number
+    except Exception:
+        filename = (
+            info["module"].replace(".", "/") + ".py"
+        )  # cannot find corresponding codeblock, use the file page instead
+
+    return "https://github.com/s3prl/s3prl/blob/%s/%s" % (tag, filename)
+
+
+# -- General configuration ---------------------------------------------------
+
+# Add any Sphinx extension module names here, as strings. They can be
+# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
+# ones.
+# add extensions
+extensions = [
+    "sphinx.ext.autodoc",
+    "sphinx.ext.napoleon",
+    "sphinx.ext.viewcode",
+    "sphinx.ext.linkcode",
+    "sphinx.ext.autosummary",
+]
+
+autosummary_generate = True
+
+html_js_files = [
+    "js/custom.js",
+]
+html_css_files = ["css/custom.css"]
+
+# Add any paths that contain templates here, relative to this directory.
+templates_path = ["_templates"]
+
+# List of patterns, relative to source directory, that match files and
+# directories to ignore when looking for source files.
+# This pattern also affects html_static_path and html_extra_path.
+exclude_patterns = []
+
+
+# -- Options for HTML output -------------------------------------------------
+
+# The theme to use for HTML and HTML Help pages.  See the documentation for
+# a list of builtin themes.
+#
+html_theme = "furo"
+
+# Add any paths that contain custom static files (such as style sheets) here,
+# relative to this directory. They are copied after the builtin static files,
+# so a file named "default.css" will overwrite the builtin "default.css".
+html_static_path = ["_static"]
+
+# Uncomment the following if you want to document __call__
+#
+# def skip(app, what, name, obj, would_skip, options):
+#     if name == "__call__":
+#         return False
+#     return would_skip
+#
+# def setup(app):
+#     app.connect("autodoc-skip-member", skip)
+
+autosummary_imported_members = True
+autosummary_ignore_module_all = False
+autodoc_member_order = "bysource"

s3prl_s3prl_main/docs/source/contribute/general.rst

@@ -0,0 +1,167 @@
+.. _general-contribution-guideline:
+
+General Guideline
+==================
+
+Thank you for considering contributing to S3PRL, we really appreciate it.
+
+However, due to the increasing difficulty of maintenance, please understand that **we might not accept the new feature**.
+
+Hence, **before submitting the implemented pull request**,
+please **submit your feature request** to the Github issue page so we can discuss about whether we want it and how to achieve it.
+
+.. warning::
+
+    If we did not go through this discussion, the pull request will not be dealt with and will be directly closed.
+
+
+Discuss
+-----------
+
+`Submit your feature request <https://github.com/s3prl/s3prl/issues/new?assignees=&labels=&template=feature_request.md&title=>`_
+on our Github Issue page to propose features and changes.
+
+Please wait for our response and do not move on to the following steps before we have a consensus on what is going to be changed.
+
+Setup
+-----------
+
+Clone the repository to **S3PRL_ROOT** and install the package
+
+.. code-block:: bash
+
+    S3PRL_ROOT="/home/leo/d/s3prl"
+    git clone https://github.com/s3prl/s3prl.git ${S3PRL_ROOT}
+    cd ${S3PRL_ROOT}
+
+    pip install -e ".[dev]"
+    # This installs the dependencies for the full functionality of S3PRL in the editable mode,
+    # including the dependencies for development like testing and building doc
+
+
+Tests
+----------
+
+Add unit tests to **${S3PRL_ROOT}/test/** to test your own new modules
+
+Verify you pass all the tests
+
+.. code-block:: bash
+
+    cd ${S3PRL_ROOT}
+    pytest
+
+
+Documentation
+-------------
+
+Make sure you write the documentation on the modules' docstring
+
+Build the documentation and make sure it looks correct
+
+.. code-block:: bash
+
+    cd ${S3PRL_ROOT}/docs
+
+    # build
+    ./rebuild_docs.sh
+
+    # launch http server
+    python3 -m http.server -d build/html/
+
+    # You can then use a browser to access the doc webpage on: YOUR_IP_OR_LOCALHOST:8000
+
+Coding-style check
+------------------
+
+Stage your changes
+
+.. code-block:: bash
+
+    git add "YOUR_CHANGED_OR_ADDED_FILES_ONLY"
+
+.. warning::
+
+    Please do not use **git add .** to add all the files under your repository.
+    If there are files not ignored by git (specified in **.gitignore**), like
+    temporary experiment result files, they will all be added into git version
+    control, which will mess out our repository.
+
+.. note::
+
+    In our **.gitignore**, there are lots of ignored files especially for *.yaml*
+    and *.sh* files. If the config files or the shell scripts are important, please
+    remember to add them forcely, for example :code::`git add -f asr.yaml`
+
+Run **pre-commit** to apply the standardized coding-style on **YOUR_CHANGED_OR_ADDED_FILES_ONLY**
+
+.. code-block:: bash
+
+    pre-commit run
+
+If the results show there are files modified by **pre-commit**, you need to re-stage
+these files following the previous step.
+
+
+Commit / Push
+-------------
+
+Commit and push the changes
+
+.. code-block:: bash
+
+    git commit -m "YOUR_COMMIT_MESSAGE"
+    git push origin "YOUR_BRANCH"
+
+
+(Optional) Test against multiple environments
+---------------------------------------------
+
+We leverage **tox** to simulate multiple envs, see the `tox configuration <https://github.com/s3prl/s3prl/blob/main/tox.ini>`_ for more information.
+Tox helps automate the pipeline of creating different virtual envs, installing differnet dependencies of S3PRL, running different testing commands.
+Our Github Action CI also relies on tox, hence you can debug the CI error locally with tox.
+
+Before using tox, make sure your cli can launch the following python versions. Usually, this can be achieved via `pyenv <https://github.com/pyenv/pyenv>`_
+
+- python3.7
+- python3.8
+- python3.9
+- python3.10
+
+List all the available environments. An environment means a pre-defined routine of packaging S3PRL, installing S3PRL, installing specific dependencies,
+test specific commands. See `tox configuration <https://github.com/s3prl/s3prl/blob/main/tox.ini>`_ for more information.
+
+.. code-block:: bash
+
+    tox -l
+
+Suppose there is an environment named :code:`all_upstream-py38-audio0.12.1`, you can also test against this specific env:
+
+.. code-block:: bash
+
+    tox -e all_upstream-py38-audio0.12.1
+
+Test all environments. This simulate the environments you will meet on the Github Action CI
+
+.. code-block:: bash
+
+    tox
+
+
+Send a pull request
+-------------------
+
+Verify your codes are in the proper format
+
+.. code-block:: bash
+
+    ./ci/format.sh --check
+    # If this fails, simply remove --check to do the actual formatting
+
+Make sure you add test cases and your change pass the tests
+
+.. code-block:: bash
+
+    pytest
+
+Send a pull request on GitHub

s3prl_s3prl_main/docs/source/contribute/private.rst

@@ -0,0 +1,104 @@
+Internal S3PRL Development
+==========================
+
+Write code
+----------
+
+1.  Make sure you have access to `s3prl/s3prl-private <https://github.com/s3prl/s3prl-private/>`_
+
+2.  Clone the repository to **S3PRL_ROOT** and install the package
+
+    .. code-block:: bash
+
+        git clone s3prl/s3prl-private "S3PRL_ROOT"
+        cd "S3PRL_ROOT"
+        pip install -e ".[dev]"  # This installs the dependencies for the full functionality of S3PRL
+
+3.  Write code into the packages listed in **S3PRL_ROOT/valid_paths.txt**
+
+Unit tests
+----------
+
+4.  Add unit tests to **S3PRL_ROOT/test/** to test your own new modules
+
+5.  Verify you pass all the tests
+
+    .. code-block:: bash
+
+        cd "S3PRL_ROOT"
+        pytest
+
+Documentation
+-------------
+
+6.  Make sure you write the documentation on the modules' docstring
+
+7.  Build the documentation and make sure it looks correct
+
+    .. code-block:: bash
+
+        cd "S3PRL_ROOT"/docs
+
+        # build
+        ./rebuild_docs.sh
+
+        # launch http server
+        python3 -m http.server -d build/html/
+
+        # You can then use a browser to access the doc webpage on: YOUR_IP_OR_LOCALHOST:8000
+
+8.  OK now your new changes are ready to be commit
+
+Coding-style check
+------------------
+
+9.  Stage your changes
+
+    .. code-block:: bash
+
+        git add "YOUR_CHANGED_OR_ADDED_FILES_ONLY"
+
+    .. warning::
+
+        Please do not use **git add .** to add all the files under your repository.
+        If there are files not ignored by git (specified in **.gitignore**), like
+        temporary experiment result files, they will all be added into git version
+        control, which will mess out our repository.
+
+10. Run **pre-commit** to apply the standardized coding-style on **YOUR_CHANGED_OR_ADDED_FILES_ONLY**
+
+    .. code-block:: bash
+
+        pre-commit run
+
+    If the results show there are files modified by **pre-commit**, you need to re-stage
+    these files following step 9.
+
+Commit / Push
+-------------
+
+12. Commit and push the changes
+
+    .. code-block:: bash
+
+        git commit -m "YOUR_COMMIT_MESSAGE"
+        git push origin "YOUR_BRANCH"
+
+Send a pull request
+-------------------
+
+Only do this when you are ready to merge your branch. Since once you send a pull request,
+every newly pushed commit will cause GitHub to run CI, but we have a limited number of
+runnable CI per month, regularized by GitHub. Hence, you should do this only after the
+branch is ready.
+
+13. Verify you can pass the CI locally
+
+    .. code-block:: bash
+
+        ./ci/format.sh --check
+        # If this fails, simply remove --check to do the actual formatting
+
+        pytest
+
+14. Send a pull request on GitHub

s3prl_s3prl_main/docs/source/contribute/public.rst

@@ -0,0 +1,29 @@
+Contribute to S3PRL
+===================
+
+Thank you for considering contributing to S3PRL! We really appreciate it!
+
+**Please discuss with us on the Github issue page about your "feature request"** before submitting
+the implemented pull request, and we will discuss about the following things:
+
+1. Whether we want to add it in S3PRL
+2. If yes, how should we achieve this feature
+3. How can we collaborate on the implementation
+4. Proceed with the actual implementation and the pull request
+
+.. warning::
+
+    If we did not go through this discussion, the pull request will not be dealt with and will be directly closed.
+
+.. note::
+
+    S3PRL has grown too large and become too hard to maintain.
+    In principle, we plan to keep only a few essential functionality which is guaranteed to work,
+    instead of continually adding lots of tasks, which can be easily added but hard to keep people
+    always there to maintain all the functions in the long term.
+    Hence, *we welcome all the bug fixes or slight improvement on existing functions,
+    while by default we do not accept pull requests containing lots of changes or an entirely new task*,
+    unless our core maintainers feel it is maintainable from our side.
+    That is, we do not expect outside collaborators to maintain their contribution for a long time,
+    which is the usual case in our experience. We wish to keep all the functionality of S3PRL work in the long term,
+    and thanks for your kind support and understanding.

s3prl_s3prl_main/docs/source/contribute/upstream.rst

@@ -0,0 +1,100 @@
+Adding New Upstream
+====================
+
+
+Discuss
+---------
+
+Please make sure that you already go through :doc:`./general`.
+Again, please make sure we have consensus on the new feature request.
+The best and the most transparent way is to
+`submit your feature request <https://github.com/s3prl/s3prl/issues/new?assignees=&labels=&template=feature_request.md&title=>`_.
+
+
+Copy from the template
+-----------------------
+
+To add new upstream, you can start with an `example <https://github.com/s3prl/s3prl/tree/main/s3prl/upstream/example>`_
+Suppose your new upstream called :code:`my_awesome_upstream`, the simplest way to start will be the following:
+
+1.
+
+.. code-block:: bash
+
+    cd ${S3PRL_ROOT}
+    cp -r s3prl/upstream/example/ s3prl/upstream/my_awesome_upstream
+
+2. In :code:`s3prl/upstream/my_awesome_upstream/hubconf.py`, change :code:`customized_upstream` to :code:`my_entry_1`
+3. In :code:`s3prl/hub.py`, add :code:`from s3prl.upstream.my_awesome_upstream.hubconf import *`
+
+4.
+
+.. code-block:: bash
+
+    python3 utility/extract_feat.py my_entry_1 sample_hidden_states
+    # this script extract hidden states from an upstream entry to the "sample_hidden_states" folder
+
+This will extract the hidden states from this :code:`my_entry_1` entry.
+The default content in :code:`s3prl/upstream/example/` always works, so you can simply edit the files
+inside the new :code:`s3prl/upstream/my_awesome_upstream` folder to enable your new upstream.
+
+
+Implement
+----------
+
+The folder is in the following structure:
+
+.. code-block:: bash
+
+    my_awesome_upstream
+    |
+     ---- expert.py
+    |
+     ---- hubconf.py
+
+In principle, :code:`hubconf.py` serves as the URL registry, where each callable function is an entry specifying
+the source of the checkpoint, while the :code:`expert.py` serves as the wrapper of your model definition to fit
+with our upstream API.
+
+During your implementation, please try to remove as many package dependencies as possible, since the upstream
+functionality is our core feature, and should have minimal dependencies to be maintainable.
+
+
+Tests
+-------
+
+After you implementation, please make sure all your entries can pass `the tests <https://github.com/s3prl/s3prl/blob/8eac602117003e2bb5cdb7a4d0e94cc9975fd4f2/test/test_upstream.py#L194-L250>`_
+The :code:`test_upstream_with_extracted` test case requires you to pre-extract the expected hidden states via:
+
+.. code-block:: bash
+
+    python3 utility/extract_feat.py my_awesome_upstream ./sample_hidden_states
+
+That is, the test case expects there will be a :code:`my_awesome_upstream.pt` in the :code:`sample_hidden_states` folder.
+
+All the existing sampled hidden states are hosted at a `Huggingface Dataset Repo <https://huggingface.co/datasets/s3prl/sample_hidden_states/tree/main>`_,
+and we expect you to clone (by :code:`git lfs`) this :code:`sample_hidden_states` repo and add the sampled hidden states for your new entries.
+
+To make changes to this hidden states repo, please follow the steps `here <https://huggingface.co/datasets/s3prl/sample_hidden_states/discussions>`_
+to create a pull request, so that our core maintainer can sync the hidden states extracted by you.
+
+In conclusion, to add new upstream one needs to make two pull requests:
+
+- To https://github.com/s3prl/s3prl/pulls
+- To https://huggingface.co/datasets/s3prl/sample_hidden_states/tree/main
+
+
+.. note::
+
+    In fact, due to the huge time cost, most of the upstreams in S3PRL will not be tested in Github Action CI (or else it will take several hours
+    to download all the checkpoints for every PRs). However, our core maintainers will still clone the repository and run tox locally to make sure
+    everything works fine, and there is a `tox environment <https://github.com/s3prl/s3prl/blob/8eac602117003e2bb5cdb7a4d0e94cc9975fd4f2/tox.ini#L11>`_
+    testing all the upstreams.
+
+
+Documentation
+--------------
+
+After all the implementation, make sure your efforts are known by the users by adding documentation of your entries at
+the :doc:`../tutorial/upstream_collection` tutorial page. Also, you can add your name at the bottom of the tutorial
+page if you like.

s3prl_s3prl_main/docs/source/index.rst

@@ -0,0 +1,75 @@
+.. S3PRL documentation master file, created by
+   sphinx-quickstart on Sun May 15 15:43:39 2022.
+   You can adapt this file completely to your liking, but it should at least
+   contain the root `toctree` directive.
+
+S3PRL
+=====
+
+.. image:: https://raw.githubusercontent.com/s3prl/s3prl/master/file/S3PRL-logo.png
+
+**S3PRL** is a toolkit targeting for Self-Supervised Learning for speech processing.
+Its full name is **S**\elf-**S**\upervised **S**\peech **P**\re-training and **R**\epresentation **L**\earning.
+It supports the following three major features:
+
+* **Pre-training**
+
+   * You can train the following models from scratch:
+
+   * *Mockingjay*, *Audio ALBERT*, *TERA*, *APC*, *VQ-APC*, *NPC*, and *DistilHuBERT*
+
+* **Pre-trained models (Upstream) collection**
+
+   * Easily load most of the existing upstream models with pretrained weights in a unified I/O interface.
+   * Pretrained models are registered through torch.hub, which means you can use these models in your own project by one-line plug-and-play without depending on this toolkit's coding style.
+
+* **Downstream Evaluation**
+
+   * Utilize upstream models in lots of downstream tasks
+   * The official implementation of the `SUPERB Benchmark <https://superbbenchmark.org/>`_
+
+
+Getting Started
+---------------
+
+.. toctree::
+   :caption: Getting started
+
+   ./tutorial/installation.rst
+   ./tutorial/upstream_collection.rst
+   ./tutorial/problem.rst
+
+
+How to Contribute
+-----------------
+
+.. toctree::
+   :caption: How to Contribute
+
+   ./contribute/general.rst
+   ./contribute/upstream.rst
+
+
+API Documentation
+-----------------
+
+.. autosummary::
+   :caption: API Documentation
+   :toctree: _autosummary
+   :template: custom-module-template.rst
+   :recursive:
+
+   s3prl.nn
+   s3prl.problem
+   s3prl.task
+   s3prl.dataio
+   s3prl.metric
+   s3prl.util
+
+
+Indices and tables
+==================
+
+* :ref:`genindex`
+* :ref:`modindex`
+* :ref:`search`

s3prl_s3prl_main/docs/source/tutorial/installation.rst

@@ -0,0 +1,55 @@
+Install S3PRL
+=============
+
+Minimal installation
+--------------------
+
+This installation only enables the **S3PRL Upstream collection** function to
+keep the minimal dependency. To enable all the functions including downstream benchmarking,
+you need to follow `Full installation`_.
+
+.. code-block:: bash
+
+    pip install s3prl
+
+
+Editable installation
+---------------------
+
+Installing a package in the editable mode means when you use a imported class/function,
+the source code of the class/function is right in your cloned repository.
+So, when you modify the code inside it, the newly imported class/function will reflect
+your modification.
+
+.. code-block:: bash
+
+    git clone https://github.com/s3prl/s3prl.git
+    cd s3prl
+    pip install -e .
+
+
+Full installation
+------------------
+
+Install all the dependencies to enable all the S3PRL functions. However, there are **LOTS**
+of dependencies.
+
+.. code-block:: bash
+
+    pip install s3prl[all]
+
+    # editable
+    pip install ".[all]"
+
+
+Development installation
+-------------------------
+
+Install dependencies of full installation and extra packages for development,
+including **pytest** for unit-testing and **sphinx** for documentation.
+
+Usually, you will use this installation variant only in editable mode
+
+.. code-block:: bash
+
+    pip install ".[dev]"

s3prl_s3prl_main/docs/source/tutorial/problem.rst

@@ -0,0 +1,122 @@
+Use Problem module to run customizable recipes
+=======================================================
+
+The :obj:`s3prl.problem` module provides customizable recipes in pure python (almost).
+See :obj:`s3prl.problem` for all the recipes ready to be ran.
+
+
+Usage 1. Import and run on Colab
+--------------------------------
+
+All the problem class follows the same usage
+
+    >>> import torch
+    >>> from s3prl.problem import SuperbASR
+    ...
+    >>> problem = SuperbASR()
+    >>> config = problem.default_config()
+    >>> print(config)
+    ...
+    >>> # See the config for the '???' required fields and fill them
+    >>> config["target_dir"] = "result/asr_exp"
+    >>> config["prepare_data"]["dataset_root"] = "/corpus/LibriSpeech/"
+    ...
+    >>> problem.run(**config)
+
+
+Usage 2. Run & configure from CLI
+-----------------------------------
+
+If you want to directly run from command-line, write a python script (:code:`asr.py`) as follow:
+
+.. code-block::
+
+    # This is asr.py
+
+    from s3prl.problem import SuperbASR
+    SuperbASR().main()
+
+Then, run the command below:
+
+    >>> # Note that the main function supports overridding a field in the config by:
+    >>> #   --{field_name} {value}
+    >>> #   --{outer_field_name}.{inner_field_name} {value}
+    ...
+    >>> python3 asr.py --target_dir result/asr_exp --prepare_data.dataset_root /corpus/LibriSpeech/
+
+
+Usage 3. Run & configure with the unified :obj:`s3prl-main`
+-----------------------------------------------------------
+
+However, this means that for every problem you still need to create a file.
+Hence, we provide an easy helper supporting all the problems in :obj:`s3prl.problem`:
+
+    >>> python3 -m s3prl.main SuperbASR --target_dir result/asr_exp --prepare_data.dataset_root /corpus/LibriSpeech/
+
+or use our CLI entry: :code:`s3prl-main`
+
+    >>> s3prl-main SuperbASR --target_dir result/asr_exp --prepare_data.dataset_root /corpus/LibriSpeech/
+
+Customization
+-------------
+
+The core feature of the :obj:`s3prl.problem` module is customization.
+You can easily change the corpus, change the SSL upstream model, change the downstream model,
+optimizer, scheduler... etc, which can all be freely defined by you!
+
+We demonstrate how to change the corpus and the downstream model in the following :code:`new_asr.py`:
+
+.. code-block:: python
+
+    # This is new_asr.py
+
+    import torch
+    import pandas as pd
+    from s3prl.problem import SuperbASR
+
+
+    class LowResourceLinearSuperbASR(SuperbASR):
+        def prepare_data(
+            self, prepare_data: dict, target_dir: str, cache_dir: str, get_path_only=False
+        ):
+            train_path, valid_path, test_paths = super().prepare_data(
+                prepare_data, target_dir, cache_dir, get_path_only
+            )
+
+            # Take only the first 100 utterances for training
+            df = pd.read_csv(train_path)
+            df = df.iloc[:100]
+            df.to_csv(train_path, index=False)
+
+            return train_path, valid_path, test_paths
+
+        def build_downstream(
+            self,
+            build_downstream: dict,
+            downstream_input_size: int,
+            downstream_output_size: int,
+            downstream_input_stride: int,
+        ):
+            class Model(torch.nn.Module):
+                def __init__(self, input_size, output_size) -> None:
+                    super().__init__()
+                    self.linear = torch.nn.Linear(input_size, output_size)
+
+                def forward(self, x, x_len):
+                    return self.linear(x), x_len
+
+            return Model(downstream_input_size, downstream_output_size)
+
+
+    if __name__ == "__main__":
+        LowResourceLinearSuperbASR().main()
+
+
+By subclassing :obj:`SuperbASR`, we create a new problem called :code:`LowResourceLinearSuperbASR` by
+overridding the :code:`prepare_data` and :code:`build_downstream` methods. After this simple modification,
+now the :code:`LowResourceLinearSuperbASR` works exactly the same as :code:`SuperbASR` while with two slight
+setting changes, and then you can follow the first two usages introduced above to launch this new class.
+
+For example:
+
+    >>> python3 new_asr.py --target_dir result/new_asr_exp --prepare_data.dataset_root /corpus/LibriSpeech/

s3prl_s3prl_main/docs/source/tutorial/upstream_collection.rst

@@ -0,0 +1,1457 @@
+:tocdepth: 2
+
+S3PRL Upstream Collection
+=======================================
+
+We collect almost all the existing SSL pre-trained models in S3PRL,
+so you can import and use them easily in an unified I/O interface.
+
+:obj:`s3prl.nn.upstream.S3PRLUpstream` is an easy interface to retrieve all the self-supervised learning (SSL) pre-trained models
+available in S3PRL. the :code:`name` argument for :obj:`s3prl.nn.upstream.S3PRLUpstream` specifies the checkpoint,
+and then the pre-trained models in this checkpoint will be automatically constructed and
+initialized.
+
+Here is an example on how to get a hubert model and its representation using the :code:`name='hubert'`:
+
+.. code-block:: python
+
+    import torch
+    from s3prl.nn import S3PRLUpstream
+
+    model = S3PRLUpstream("hubert")
+    model.eval()
+
+    with torch.no_grad():
+        wavs = torch.randn(2, 16000 * 2)
+        wavs_len = torch.LongTensor([16000 * 1, 16000 * 2])
+        all_hs, all_hs_len = model(wavs, wavs_len)
+
+    for hs, hs_len in zip(all_hs, all_hs_len):
+        assert isinstance(hs, torch.FloatTensor)
+        assert isinstance(hs_len, torch.LongTensor)
+
+        batch_size, max_seq_len, hidden_size = hs.shape
+        assert hs_len.dim() == 1
+
+.. tip::
+
+    For each SSL learning method, like wav2vec 2.0, there are several checkpoint variants, trained by
+    different amount of unlabeled data, or different model sizes. Hence there are also various
+    :code:`name` to retrieve these different models.
+
+    Like, the HuBERT method has "hubert" and "hubert_large_ll60k" different names for different
+    checkpoint variants.
+
+.. tip::
+
+    Some SSL pre-trained models' entries can be further configured by a :code:`extra_conf` dictionary.
+    See :obj:`s3prl.nn.S3PRLUpstream`. You can find the valid :code:`extra_conf` options in each SSL
+    model category. If not documented, by default it does not support any :code:`extra_conf`.
+
+The following includes the model and checkpoint information for each :code:`name`, including the releasing date,
+paper, citation, model architecture, pre-training data, criterion, and their source code. The format follows:
+
+
+
+SSL Method
+--------------------------------------------------------
+`Paper full title with arxiv link <https://arxiv.org/>`_
+
+.. code-block:: bash
+
+    @article{citation-block,
+        title={Paper Title},
+        author={Authors},
+        year={2020},
+        month={May}
+    }
+
+The information shared across checkpoint variants.
+
+name1
+~~~~~~~~~~~~~~~~~~~
+
+The detailed specific information for this checkpoint variant (:code:`name=name1`)
+
+name2
+~~~~~~~~~~~~~~~~~~~
+
+The detailed specific information for this checkpoint variant (:code:`name=name2`)
+
+
+
+Mockingjay
+--------------------------------------------------------
+`Mockingjay: Unsupervised Speech Representation Learning with Deep Bidirectional Transformer Encoders <https://arxiv.org/abs/1910.12638>`_
+
+.. code-block:: bash
+
+    @article{mockingjay,
+        title={Mockingjay: Unsupervised Speech Representation Learning with Deep Bidirectional Transformer Encoders},
+        ISBN={9781509066315},
+        url={http://dx.doi.org/10.1109/ICASSP40776.2020.9054458},
+        DOI={10.1109/icassp40776.2020.9054458},
+        journal={ICASSP 2020 - 2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)},
+        publisher={IEEE},
+        author={Liu, Andy T. and Yang, Shu-wen and Chi, Po-Han and Hsu, Po-chun and Lee, Hung-yi},
+        year={2020},
+        month={May}
+    }
+
+Mockingjay is a BERT on Spectrogram, with 12-layers of transformer encoders in the paper.
+
+
+mockingjay
+~~~~~~~~~~~~~~~~
+
+This is alias for `mockingjay_origin`_
+
+mockingjay_origin
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+This is alias for `mockingjay_logMelLinearLarge_T_AdamW_b32_500k_360hr_drop1`_
+
+mockingjay_100hr
+~~~~~~~~~~~~~~~~
+
+This is alias for `mockingjay_logMelBase_T_AdamW_b32_200k_100hr`_
+
+mockingjay_960hr
+~~~~~~~~~~~~~~~~
+
+This is alias for `mockingjay_logMelBase_T_AdamW_b32_1m_960hr_drop1`_
+
+mockingjay_logMelBase_T_AdamW_b32_200k_100hr
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+- Feature: 80-dim log Mel
+- Alteration: time
+- Optimizer: AdamW
+- Batch size: 32
+- Total steps: 200k
+- Unlabled Speech: LibriSpeech 100hr
+
+mockingjay_logMelLinearLarge_T_AdamW_b32_500k_360hr_drop1
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+- Feature: 80-dim log Mel (input) / 201-dim Linear (target)
+- Alteration: time
+- Optimizer: AdamW
+- Batch size: 32
+- Total steps: 500k
+- Unlabled Speech: LibriSpeech 360hr
+
+mockingjay_logMelBase_T_AdamW_b32_1m_960hr
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+- Feature: 80-dim log Mel
+- Alteration: time
+- Optimizer: AdamW
+- Batch size: 32
+- Total steps: 1M
+- Unlabled Speech: LibriSpeech 960hr
+
+mockingjay_logMelBase_T_AdamW_b32_1m_960hr_drop1
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+- Feature: 80-dim log Mel
+- Alteration: time
+- Optimizer: AdamW
+- Batch size: 32
+- Total steps: 1M
+- Unlabled Speech: LibriSpeech 960hr
+- Differences: Dropout of 0.1 (instead of 0.3)
+
+
+mockingjay_logMelBase_T_AdamW_b32_1m_960hr_seq3k
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+- Feature: 80-dim log Mel
+- Alteration: time
+- Optimizer: AdamW
+- Batch size: 32
+- Total steps: 1M
+- Unlabled Speech: LibriSpeech 960hr
+- Differences: sequence length of 3k (instead of 1.5k)
+
+
+
+TERA
+--------------------------------------------------------
+`TERA: Self-Supervised Learning of Transformer Encoder Representation for Speech <https://arxiv.org/abs/2007.06028>`_
+
+.. code-block:: bash
+
+    @misc{tera,
+        title={TERA: Self-Supervised Learning of Transformer Encoder Representation for Speech},
+        author={Andy T. Liu and Shang-Wen Li and Hung-yi Lee},
+        year={2020},
+        eprint={2007.06028},
+        archivePrefix={arXiv},
+        primaryClass={eess.AS}
+    }
+
+
+tera
+~~~~~~~~~~~~~~~~
+
+This is alias for `tera_960hr`_
+
+tera_100hr
+~~~~~~~~~~~~~~~~~~
+
+This is alias for `tera_logMelBase_T_F_M_AdamW_b32_200k_100hr`_
+
+tera_960hr
+~~~~~~~~~~~~~~~~~~~
+
+This is alias for `tera_logMelBase_T_F_M_AdamW_b32_1m_960hr_drop1`_
+
+tera_logMelBase_T_F_AdamW_b32_200k_100hr
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+- Feature: 80-dim log Mel
+- Alteration: time + freq
+- Optimizer: AdamW
+- Batch size: 32
+- Total steps: 200k
+- Unlabled Speech: LibriSpeech 100hr
+
+tera_logMelBase_T_F_M_AdamW_b32_200k_100hr
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+- Feature: 80-dim log Mel
+- Alteration: time + freq + mag
+- Optimizer: AdamW
+- Batch size: 32
+- Total steps: 200k
+- Unlabled Speech: LibriSpeech 100hr
+
+tera_logMelBase_T_F_AdamW_b32_1m_960hr
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+- Feature: 80-dim log Mel
+- Alteration: time + freq
+- Optimizer: AdamW
+- Batch size: 32
+- Total steps: 1M
+- Unlabled Speech: LibriSpeech 960hr
+
+tera_logMelBase_T_F_AdamW_b32_1m_960hr_drop1
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+- Feature: 80-dim log Mel
+- Alteration: time + freq
+- Optimizer: AdamW
+- Batch size: 32
+- Total steps: 1M
+- Unlabled Speech: LibriSpeech 960hr
+- Differences: Dropout of 0.1 (instead of 0.3)
+
+tera_logMelBase_T_F_AdamW_b32_1m_960hr_seq3k
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+- Feature: 80-dim log Mel
+- Alteration: time + freq
+- Optimizer: AdamW
+- Batch size: 32
+- Total steps: 1M
+- Unlabled Speech: LibriSpeech 960hr
+- Differences: sequence length of 3k (instead of 1.5k)
+
+tera_logMelBase_T_F_M_AdamW_b32_1m_960hr_drop1
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+- Feature: 80-dim log Mel
+- Alteration: time + freq + mag
+- Optimizer: AdamW
+- Batch size: 32
+- Total steps: 1M
+- Unlabled Speech: 960hr
+- Differences: Dropout of 0.1 (instead of 0.3)
+
+tera_fbankBase_T_F_AdamW_b32_200k_100hr
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+- Feature: 240-dim fbank
+- Alteration: time + freq
+- Optimizer: AdamW
+- Batch size: 32
+- Total steps: 200k
+- Unlabled Speech: LibriSpeech 100hr
+
+
+
+Audio ALBERT
+--------------------------------------------------------
+`Audio ALBERT: A Lite BERT for Self-supervised Learning of Audio Representation <https://arxiv.org/abs/2007.06028>`_
+
+.. code-block:: bash
+
+    @inproceedings{chi2021audio,
+        title={Audio albert: A lite bert for self-supervised learning of audio representation},
+        author={Chi, Po-Han and Chung, Pei-Hung and Wu, Tsung-Han and Hsieh, Chun-Cheng and Chen, Yen-Hao and Li, Shang-Wen and Lee, Hung-yi},
+        booktitle={2021 IEEE Spoken Language Technology Workshop (SLT)},
+        pages={344--350},
+        year={2021},
+        organization={IEEE}
+    }
+
+
+audio_albert
+~~~~~~~~~~~~~~~~
+
+This is alias of `audio_albert_960hr`_
+
+
+audio_albert_960hr
+~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+This is alias of `audio_albert_logMelBase_T_share_AdamW_b32_1m_960hr_drop1`_
+
+
+audio_albert_logMelBase_T_share_AdamW_b32_1m_960hr_drop1
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+- Feature: 80-dim log Mel
+- Alteration: time
+- Optimizer: AdamW
+- Batch size: 32
+- Total steps: 1M
+- Unlabled Speech: LibriSpeech 960hr
+
+
+
+APC
+--------------------------------------------------------
+`An Unsupervised Autoregressive Model for Speech Representation Learning <https://arxiv.org/abs/1904.03240>`_
+
+.. code-block:: bash
+
+    @inproceedings{chung2019unsupervised,
+        title = {An unsupervised autoregressive model for speech representation learning},
+        author = {Chung, Yu-An and Hsu, Wei-Ning and Tang, Hao and Glass, James},
+        booktitle = {Interspeech},
+        year = {2019}
+    }
+
+
+apc
+~~~~~~~~~~~~~~~~
+
+This is alias of `apc_360hr`_
+
+
+apc_360hr
+~~~~~~~~~~~~~~~~~~
+
+- Unlabled Speech: LibriSpeech 360hr
+
+
+apc_960hr
+~~~~~~~~~~~~~~~~~~
+
+- Unlabled Speech: LibriSpeech 960hr
+
+
+
+VQ-APC
+--------------------------------------------------------
+`Vector-Quantized Autoregressive Predictive Coding <https://arxiv.org/abs/2005.08392>`_
+
+.. code-block:: bash
+
+    @inproceedings{chung2020vqapc,
+        title = {Vector-quantized autoregressive predictive coding},
+        autohor = {Chung, Yu-An and Tang, Hao and Glass, James},
+        booktitle = {Interspeech},
+        year = {2020}
+    }
+
+vq_apc
+~~~~~~~~~~~~~~~~
+
+This is alias of `vq_apc_360hr`_
+
+
+vq_apc_360hr
+~~~~~~~~~~~~~~~~
+
+- Unlabled Speech: LibriSpeech 360hr
+
+
+vq_apc_960hr
+~~~~~~~~~~~~~~~~~
+
+- Unlabled Speech: LibriSpeech 960hr
+
+
+
+NPC
+--------------------------------------------------------
+`Non-Autoregressive Predictive Coding for Learning Speech Representations from Local Dependencies <https://arxiv.org/abs/2011.00406>`_
+
+.. code-block:: bash
+
+    @article{liu2020nonautoregressive,
+        title   = {Non-Autoregressive Predictive Coding for Learning Speech Representations from Local Dependencies},
+        author  = {Liu, Alexander and Chung, Yu-An and Glass, James},
+        journal = {arXiv preprint arXiv:2011.00406},
+        year    = {2020}
+    }
+
+
+npc
+~~~~~~~~~~~~~~~~
+
+This is alias of `npc_360hr`_
+
+
+npc_360hr
+~~~~~~~~~~~~~~~~~~
+
+- Unlabled Speech: LibriSpeech 360hr
+
+
+npc_960hr
+~~~~~~~~~~~~~~~~~~~~
+
+- Unlabled Speech: LibriSpeech 960hr
+
+
+
+PASE+
+--------------------------------------------------------
+`Multi-task self-supervised learning for Robust Speech Recognition <https://arxiv.org/abs/2001.09239>`_
+
+.. code-block:: bash
+
+    @inproceedings{ravanelli2020multi,
+        title={Multi-task self-supervised learning for robust speech recognition},
+        author={Ravanelli, Mirco and Zhong, Jianyuan and Pascual, Santiago and Swietojanski, Pawel and Monteiro, Joao and Trmal, Jan and Bengio, Yoshua},
+        booktitle={ICASSP 2020-2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)},
+        pages={6989--6993},
+        year={2020},
+        organization={IEEE}
+    }
+
+.. hint::
+
+    To use PASE models, there are many extra dependencies required to install.
+    Please follow the below installation instruction:
+
+    .. code-block:: bash
+
+        pip install -r https://raw.githubusercontent.com/s3prl/s3prl/master/s3prl/upstream/pase/requirements.txt
+
+
+pase_plus
+~~~~~~~~~~~~~~~~
+
+- Unlabled Speech: LibriSpeech 50hr
+
+
+
+Modified CPC
+--------------------------------------------------------
+`Unsupervised pretraining transfers well across languages <https://arxiv.org/abs/2002.02848>`_
+
+.. code-block:: bash
+
+    @inproceedings{riviere2020unsupervised,
+        title={Unsupervised pretraining transfers well across languages},
+        author={Riviere, Morgane and Joulin, Armand and Mazar{\'e}, Pierre-Emmanuel and Dupoux, Emmanuel},
+        booktitle={ICASSP 2020-2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)},
+        pages={7414--7418},
+        year={2020},
+        organization={IEEE}
+    }
+
+.. note::
+
+    This is a slightly improved version on the original CPC by DeepMind. To cite the DeepMind version:
+
+    .. code-block:: bash
+
+        @article{oord2018representation,
+            title={Representation learning with contrastive predictive coding},
+            author={Oord, Aaron van den and Li, Yazhe and Vinyals, Oriol},
+            journal={arXiv preprint arXiv:1807.03748},
+            year={2018}
+        }
+
+
+modified_cpc
+~~~~~~~~~~~~~~~~
+
+- Unlabled Speech: LibriLight 60k hours
+
+
+
+DeCoAR
+--------------------------------------------------------
+`Deep contextualized acoustic representations for semi-supervised speech recognition <https://arxiv.org/abs/1912.01679>`_
+
+.. code-block:: bash
+
+    @inproceedings{ling2020deep,
+        title={Deep contextualized acoustic representations for semi-supervised speech recognition},
+        author={Ling, Shaoshi and Liu, Yuzong and Salazar, Julian and Kirchhoff, Katrin},
+        booktitle={ICASSP 2020-2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)},
+        pages={6429--6433},
+        year={2020},
+        organization={IEEE}
+    }
+
+
+decoar_layers
+~~~~~~~~~~~~~~~~
+
+- Unlabled Speech: LibriSpeech 960hr
+
+
+DeCoAR 2.0
+--------------------------------------------------------
+`DeCoAR 2.0: Deep Contextualized Acoustic Representations with Vector Quantization <https://arxiv.org/abs/2012.06659>`_
+
+.. code-block:: bash
+
+    @misc{ling2020decoar,
+        title={DeCoAR 2.0: Deep Contextualized Acoustic Representations with Vector Quantization}, 
+        author={Shaoshi Ling and Yuzong Liu},
+        year={2020},
+        eprint={2012.06659},
+        archivePrefix={arXiv},
+        primaryClass={eess.AS}
+    }
+
+
+decoar2
+~~~~~~~~~~~~~~~~~~~~~
+
+- Unlabled Speech: LibriSpeech 960hr
+
+
+
+wav2vec
+--------------------------------------------------
+`wav2vec: Unsupervised Pre-Training for Speech Recognition <https://arxiv.org/abs/1904.05862>`_
+
+.. code-block:: bash
+
+    @article{schneider2019wav2vec,
+        title={wav2vec: Unsupervised Pre-Training for Speech Recognition},
+        author={Schneider, Steffen and Baevski, Alexei and Collobert, Ronan and Auli, Michael},
+        journal={Proc. Interspeech 2019},
+        pages={3465--3469},
+        year={2019}
+    }
+
+
+wav2vec
+~~~~~~~~~~~
+
+This is alias of `wav2vec_large`_
+
+
+wav2vec_large
+~~~~~~~~~~~~~~~
+
+This is the official wav2vec model from fairseq.
+
+- Unlabled Speech: LibriSpeech 960hr
+
+
+vq-wav2vec
+--------------------------------------------------
+`vq-wav2vec: Self-supervised learning of discrete speech representations <https://arxiv.org/abs/1910.05453>`_
+
+.. code-block:: bash
+
+    @inproceedings{baevski2019vq,
+        title={vq-wav2vec: Self-Supervised Learning of Discrete Speech Representations},
+        author={Baevski, Alexei and Schneider, Steffen and Auli, Michael},
+        booktitle={International Conference on Learning Representations},
+        year={2019}
+    }
+
+.. note::
+
+    We only take the Conv encoders' hidden_states for vq-wav2vec in this SSL method category.
+    If you wish to consider the BERT model after ths Conv encoders, please refer to `Discrete BERT`_.
+
+vq_wav2vec
+~~~~~~~~~~~
+
+This is alias of `vq_wav2vec_gumbel`_
+
+
+vq_wav2vec_gumbel
+~~~~~~~~~~~~~~~~~~~~
+
+This is the official vq-wav2vec model from fairseq.
+This model uses gumbel-softmax as the quantization technique
+
+- Unlabled Speech: LibriSpeech 960hr
+
+
+vq_wav2vec_kmeans
+~~~~~~~~~~~~~~~~~~~~~
+
+This is the official vq-wav2vec model from fairseq.
+This model uses K-means as the quantization technique
+
+
+Discrete BERT
+--------------------------------------------------
+`vq-wav2vec: Self-supervised learning of discrete speech representations <https://arxiv.org/abs/1910.05453>`_
+
+.. code-block:: bash
+
+    @inproceedings{baevski2019vq,
+        title={vq-wav2vec: Self-Supervised Learning of Discrete Speech Representations},
+        author={Baevski, Alexei and Schneider, Steffen and Auli, Michael},
+        booktitle={International Conference on Learning Representations},
+        year={2019}
+    }
+
+This method takes the Conv feature encoder's output, quantize it into token ids, and feed the
+tokens into a NLP BERT (Specifically, RoBERTa). The output hidden_states are all the hidden hidden_states
+of the NLP BERT (excluding the hidden_states in `vq-wav2vec`_)
+
+
+discretebert
+~~~~~~~~~~~~~~~~
+
+Alias of `vq_wav2vec_kmeans_roberta`_
+
+
+vq_wav2vec_kmeans_roberta
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+This model uses `vq_wav2vec_kmeans`_ as the frontend waveform tokenizer. After the waveform is tokenized
+into a sequence of token ids, tokens are then fed into a RoBERTa model.
+
+
+
+wav2vec 2.0
+--------------------------------------------------
+`wav2vec 2.0: A Framework for Self-Supervised Learning of Speech Representations <https://arxiv.org/abs/2006.11477>`_
+
+.. code-block:: bash
+
+    @article{baevski2020wav2vec,
+        title={wav2vec 2.0: A framework for self-supervised learning of speech representations},
+        author={Baevski, Alexei and Zhou, Yuhao and Mohamed, Abdelrahman and Auli, Michael},
+        journal={Advances in Neural Information Processing Systems},
+        volume={33},
+        pages={12449--12460},
+        year={2020}
+    }
+
+All the entries below support the following :code:`extra_conf`:
+
+====================  ====================
+column                description
+====================  ====================
+feature_selection     (str) -
+                        if :code:`fairseq_layers` or :code:`fairseq_layers_before_residual`,
+                        extract the representation following official fairseq API.
+                        for :code:`fairseq_layers`, it is the output of each transformer
+                        encoder layer; for :code:`fairseq_layers_before_residual`, it is
+                        the output of the feedforward layer (before adding with the
+                        main residual) of each transformer encoder layer. by default
+                        this option is None, which follows the default place to extract
+                        in S3PRL.
+====================  ====================
+
+
+wav2vec2_custom
+~~~~~~~~~~~~~~~~~~~~~
+
+This entry expects you to provide the source of the checkpoint: :code:`path_or_url`, which should be
+the local path or a url of the checkpoint converted by :code:`s3prl/upstream/wav2vec2/convert.py` (
+from a regular fairseq checkpoint.)
+
+This entry also supports the following additional :code:`extra_conf`.
+
+====================  ====================
+column                description
+====================  ====================
+fairseq               (bool) -
+                        If True, perform the on-the-fly checkpoint conversion, so that
+                        you can directly give the fairseq checkpoint to the :code:`path_or_url`
+                        argument, either a fairseq URL or a fairseq checkpoint local path.
+====================  ====================
+
+
+hf_wav2vec2_custom
+~~~~~~~~~~~~~~~~~~~~
+
+This entry expects you to provide the source of the checkpoint: :code:`path_or_url`, which should be
+in the HuggingFace format, like :code:`facebook/wav2vec2-large-960h`
+
+
+wav2vec2
+~~~~~~~~~~~~~~~~
+
+This is the alias of `wav2vec2_base_960`_
+
+
+wav2vec2_base_960
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+This is the official wav2vec 2.0 model in fairseq
+
+- Architecture: 12-layer Transformer encoders
+- Unlabled Speech: LibriSpeech 960hr
+
+
+wav2vec2_large_960
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+- Architecture: 24-layer Transformer encoders
+- Unlabled Speech: LibriSpeech 960hr
+
+
+wav2vec2_large_ll60k
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+- Architecture: 24-layer Transformer encoders
+- Unlabled Speech: LibriLight LL60k hours
+
+
+wav2vec2_large_lv60_cv_swbd_fsh
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The Large model trained on Libri-Light 60k hours + CommonVoice + Switchboard + Fisher
+
+- Architecture: 24-layer Transformer encoders
+- Unlabeled Speech: Libri-Light 60k hours + CommonVoice + Switchboard + Fisher
+
+
+wav2vec2_conformer_relpos
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The results can be found in the Table 4 of `fairseq S2T: Fast Speech-to-Text Modeling with fairseq <https://arxiv.org/abs/2010.05171>`_.
+
+- Architecture: 24-layer Conformer encoders with relative positional encoding
+- Unlabeled Speech: LibriLight LL60k hours
+
+
+wav2vec2_conformer_rope
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The results can be found in the Table 4 of `fairseq S2T: Fast Speech-to-Text Modeling with fairseq <https://arxiv.org/abs/2010.05171>`_.
+
+- Architecture: 24-layer Conformer encoders with ROPE positional encoding
+- Unlabeled Speech: LibriLight LL60k hours
+
+
+wav2vec2_base_s2st_es_voxpopuli
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+- The wav2vec2 model from `Enhanced Direct Speech-to-Speech Translation Using Self-supervised Pre-training and Data Augmentation <https://arxiv.org/abs/2204.02967>`_,
+- released in Fairseq with the link: `https://dl.fbaipublicfiles.com/fairseq/speech_to_speech/s2st_finetuning/w2v2/es/transformer_B.pt <https://dl.fbaipublicfiles.com/fairseq/speech_to_speech/s2st_finetuning/w2v2/es/transformer_B.pt>`_
+
+
+wav2vec2_base_s2st_en_librilight
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+- The wav2vec2 model from `Enhanced Direct Speech-to-Speech Translation Using Self-supervised Pre-training and Data Augmentation <https://arxiv.org/abs/2204.02967>`_,
+- released in Fairseq with the link: `https://dl.fbaipublicfiles.com/fairseq/speech_to_speech/s2st_finetuning/w2v2/en/transformer_B.pt <https://dl.fbaipublicfiles.com/fairseq/speech_to_speech/s2st_finetuning/w2v2/en/transformer_B.pt>`_
+
+
+wav2vec2_conformer_large_s2st_es_voxpopuli
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+- The wav2vec2 model from `Enhanced Direct Speech-to-Speech Translation Using Self-supervised Pre-training and Data Augmentation <https://arxiv.org/abs/2204.02967>`_,
+- released in Fairseq with the link: `https://dl.fbaipublicfiles.com/fairseq/speech_to_speech/s2st_finetuning/w2v2/es/conformer_L.pt <https://dl.fbaipublicfiles.com/fairseq/speech_to_speech/s2st_finetuning/w2v2/es/conformer_L.pt>`_
+
+
+wav2vec2_conformer_large_s2st_en_librilight
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+- The wav2vec2 model from `Enhanced Direct Speech-to-Speech Translation Using Self-supervised Pre-training and Data Augmentation <https://arxiv.org/abs/2204.02967>`_,
+- released in Fairseq with the link: `https://dl.fbaipublicfiles.com/fairseq/speech_to_speech/s2st_finetuning/w2v2/en/conformer_L.pt <https://dl.fbaipublicfiles.com/fairseq/speech_to_speech/s2st_finetuning/w2v2/en/conformer_L.pt>`_
+
+
+xlsr_53
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The wav2vec 2.0 model trained on multilingual presented in `Unsupervised Cross-lingual Representation Learning for Speech Recognition <https://arxiv.org/abs/2006.13979>`_
+
+.. code-block:: bash
+
+    @article{conneau2020unsupervised,
+        title={Unsupervised cross-lingual representation learning for speech recognition},
+        author={Conneau, Alexis and Baevski, Alexei and Collobert, Ronan and Mohamed, Abdelrahman and Auli, Michael},
+        journal={arXiv preprint arXiv:2006.13979},
+        year={2020}
+    }
+
+
+XLS-R
+--------------------------------------------------
+`XLS-R: Self-supervised Cross-lingual Speech Representation Learning at Scale <https://arxiv.org/abs/2111.09296>`_
+
+.. code-block:: bash
+
+    @article{babu2021xls,
+        title={XLS-R: Self-supervised cross-lingual speech representation learning at scale},
+        author={Babu, Arun and Wang, Changhan and Tjandra, Andros and Lakhotia, Kushal and Xu, Qiantong and Goyal, Naman and Singh, Kritika and von Platen, Patrick and Saraf, Yatharth and Pino, Juan and others},
+        journal={arXiv preprint arXiv:2111.09296},
+        year={2021}
+    }
+
+
+xls_r_300m
+~~~~~~~~~~~~~~~~~~~~~
+
+- Unlabled Speech: 128 languages, 436K hours
+
+
+xls_r_1b
+~~~~~~~~~~~~~~~~~~~~~
+
+- Unlabled Speech: 128 languages, 436K hours
+
+
+xls_r_2b
+~~~~~~~~~~~~~~~~~~~~~
+
+- Unlabled Speech: 128 languages, 436K hours
+
+
+HuBERT
+--------------------------------------------------
+`HuBERT: Self-Supervised Speech Representation Learning by Masked Prediction of Hidden Units <https://arxiv.org/abs/2106.07447>`_
+
+.. code-block:: bash
+
+    @article{hsu2021hubert,
+        title={Hubert: Self-supervised speech representation learning by masked prediction of hidden units},
+        author={Hsu, Wei-Ning and Bolte, Benjamin and Tsai, Yao-Hung Hubert and Lakhotia, Kushal and Salakhutdinov, Ruslan and Mohamed, Abdelrahman},
+        journal={IEEE/ACM Transactions on Audio, Speech, and Language Processing},
+        volume={29},
+        pages={3451--3460},
+        year={2021},
+        publisher={IEEE}
+    }
+
+
+hubert_custom
+~~~~~~~~~~~~~~~~~~~~~
+
+This entry expects you to provide the source of the checkpoint: :code:`path_or_url`, which should be
+the local path or a url of the checkpoint converted by :code:`s3prl/upstream/hubert/convert.py` (
+from a regular fairseq checkpoint.)
+
+This entry also supports the following additional :code:`extra_conf`.
+
+====================  ====================
+column                description
+====================  ====================
+fairseq               (bool) -
+                        If True, perform the on-the-fly checkpoint conversion, so that
+                        you can directly give the fairseq checkpoint to the :code:`path_or_url`
+                        argument, either a fairseq URL or a fairseq checkpoint local path.
+====================  ====================
+
+
+hf_hubert_custom
+~~~~~~~~~~~~~~~~~~~~
+
+This entry expects you to provide the source of the checkpoint: :code:`path_or_url`, which should be
+in the HuggingFace format, like :code:`facebook/hubert-large-ll60k`
+
+
+hubert
+~~~~~~~~~~~~~~~~~~~~~
+
+This is alias of `hubert_base`_
+
+
+hubert_base
+~~~~~~~~~~~~~~~~~~~~~
+
+- Unlabled Speech: LibriSpeech 960hr
+
+
+hubert_large_ll60k
+~~~~~~~~~~~~~~~~~~~~~
+
+- Unlabled Speech: LibriLight ll60k hours
+
+
+mhubert_base_vp_en_es_fr_it3
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+- The multilingual model from `Textless Speech-to-Speech Translation on Real Data <https://arxiv.org/abs/2112.08352>`_
+
+
+ESPnetHuBERT
+----------------------
+`Reducing Barriers to Self-Supervised Learning: HuBERT Pre-training with Academic Compute <https://arxiv.org/abs/2306.06672>`_
+
+.. code-block:: bash
+
+    @inproceedings{chen23l_interspeech,
+        author={William Chen and Xuankai Chang and Yifan Peng and Zhaoheng Ni and Soumi Maiti and Shinji Watanabe},
+        title={{Reducing Barriers to Self-Supervised Learning: HuBERT Pre-training with Academic Compute}},
+        year=2023,
+        booktitle={Proc. INTERSPEECH 2023},
+        pages={4404--4408},
+        doi={10.21437/Interspeech.2023-1176}
+    }
+
+
+espnet_hubert_custom
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+This entry expects you to provide the source of the checkpoint: :code:`ckpt`, which should be
+the local path of the checkpoint pretrained from ESPnet (e.g., latest.pth).
+
+
+espnet_hubert_base_iter0
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+- Unlabeled Speech: LibriSpeech 960hr (first iteration of HuBERT pre-training)
+
+
+espnet_hubert_base_iter1
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+- Unlabeled Speech: LibriSpeech 960hr (second iteration of HuBERT pre-training)
+
+
+espnet_hubert_large_gs_ll60k
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+- Unlabeled Speech: LibriLight ll60k hours
+- Labeled Speech: GigaSpeech 10k hours (to get units)
+
+
+WavLabLM
+----------------------
+`Joint Prediction and Denoising for Large-scale Multilingual Self-supervised Learning <https://arxiv.org/abs/2309.15317>`_
+
+.. code-block:: bash
+
+    @inproceedings{chen23joint,
+        author={William Chen and Jiatong Shi and Brian Yan and Dan Berrebbi and Wangyou Zhang and Yifan Peng and Xuankai Chang and Soumi Maiti and Shinji Watanabe},
+        title={Joint Prediction and Denoising for Large-scale Multilingual Self-supervised Learning},
+        year=2023,
+        booktitle={IEEE Automatic Speech Recognition and Understanding Workshop (ASRU)},
+    }
+
+
+cvhubert
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+- Unlabeled Speech: Commonvoice V11 Multilingual Data (13.6k hours)
+- only 20ms resolution version is provided. `check huggingface  for other resolutions <https://huggingface.co/espnet/espnet_cvhubert/tree/main>`_
+
+
+wavlablm_ek_40k
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+- Unlabeled Speech: Openli110 (Combination of Commonvoice, Voxpopuli, MLS, Googlei18n, around 39k hours)
+- Initialed from hubert_large_ll60k and continue train with English based k-means from librispeech
+
+
+wavlablm_mk_40k
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+- Unlabeled Speech: Openli110 (Combination of Commonvoice, Voxpopuli, MLS, Googlei18n, around 39k hours)
+- Trained from scratch and use a multilingual k-means from the training data
+
+
+wavlablm_ms_40k
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+- Unlabeled Speech: Openli110 (Combination of Commonvoice, Voxpopuli, MLS, Googlei18n, around 39k hours)
+- Trained from scratch and use a multilingual k-means from the training data with a multi-stage training
+
+
+Multiresolution HuBERT (MR-HuBERT)
+----------------------
+`Multi-resolution HuBERT: Multi-resolution Speech Self-Supervised Learning with Masked Unit Prediction <https://openreview.net/pdf?id=kUuKFW7DIF>`_
+
+.. code-block:: bash
+
+    @inproceedings{anonymous2023multiresolution,
+        title={Multi-resolution Hu{BERT}: Multi-resolution Speech Self-Supervised Learning with Masked Unit Prediction},
+        author={Anonymous},
+        booktitle={Submitted to The Twelfth International Conference on Learning Representations},
+        year={2023},
+        url={https://openreview.net/forum?id=kUuKFW7DIF},
+        note={under review}
+    }
+
+
+multires_hubert_custom
+~~~~~~~~~~~~~~~~~~~~~
+
+This entry expects you to provide the source of the checkpoint: :code:`ckpt`, which should be
+the local path or a url of the checkpoint converted by :code:`s3prl/upstream/multires_hubert/convert.py` (
+from a regular fairseq checkpoint.)
+For more available checkpoints, please check `Fairseq official release <https://github.com/facebookresearch/fairseq/blob/main/examples/mr_hubert/README.md>`_
+Related converted checkpoints are also at `S3PRL HuggingFace Repo <https://huggingface.co/s3prl/mr_hubert>`_
+
+
+multires_hubert_base
+~~~~~~~~~~~~~~~~~~~~~
+
+- Unlabled Speech: LibriSpeech 960hr
+- K-means extracted from `hubert_base`_
+
+
+multires_hubert_large
+~~~~~~~~~~~~~~~~~~~~~
+
+- Unlabeled Speech: LibriLight 60khr
+- K-means extracted from `hubert_base`_
+
+
+multires_hubert_multilingual_base
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+- Unlabeled Speech: Voxpopuli 100khr
+- K-means extracted from `hubert_base`_
+
+
+multires_hubert_multilingual_large400k
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+- Unlabeled Speech: Voxpopuli 100khr
+- K-means extracted from `hubert_base`_
+- Training steps 400k
+
+
+multires_hubert_multilingual_large600k
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+- Unlabeled Speech: Voxpopuli 100khr
+- K-means extracted from `hubert_base`_
+- Training steps 600k
+
+
+DistilHuBERT
+----------------------
+`DistilHuBERT: Speech Representation Learning by Layer-wise Distillation of Hidden-unit BERT <https://arxiv.org/abs/2110.01900>`_
+
+.. code-block:: bash
+
+    @inproceedings{chang2022distilhubert,
+        title={DistilHuBERT: Speech representation learning by layer-wise distillation of hidden-unit BERT},
+        author={Chang, Heng-Jui and Yang, Shu-wen and Lee, Hung-yi},
+        booktitle={ICASSP 2022-2022 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)},
+        pages={7087--7091},
+        year={2022},
+        organization={IEEE}
+    }
+
+
+distilhubert
+~~~~~~~~~~~~~~~~~~~~~
+
+Alias of `distilhubert_base`_
+
+
+distilhubert_base
+~~~~~~~~~~~~~~~~~~~~~
+
+- Teacher: `hubert_base`_
+- Unlabled Speech: LibriSpeech 960hr
+
+
+HuBERT-MGR
+--------------------------------------------------
+`Improving Distortion Robustness of Self-supervised Speech Processing Tasks with Domain Adaptation <https://arxiv.org/abs/2203.16104>`_
+
+.. code-block:: bash
+
+    @article{huang2022improving,
+        title={Improving Distortion Robustness of Self-supervised Speech Processing Tasks with Domain Adaptation},
+        author={Huang, Kuan Po and Fu, Yu-Kuan and Zhang, Yu and Lee, Hung-yi},
+        journal={arXiv preprint arXiv:2203.16104},
+        year={2022}
+    }
+
+
+hubert_base_robust_mgr
+~~~~~~~~~~~~~~~~~~~~~~~
+
+- Unlabled Speech: LibriSpeech 960hr
+- Augmentation: MUSAN, gaussian, reverberation
+
+
+Unispeech-SAT
+--------------------------------------------------
+`Unispeech-sat: Universal speech representation learning with speaker aware pre-training <https://arxiv.org/abs/2110.05752>`_
+
+.. code-block:: bash
+
+    @inproceedings{chen2022unispeech,
+        title={Unispeech-sat: Universal speech representation learning with speaker aware pre-training},
+        author={Chen, Sanyuan and Wu, Yu and Wang, Chengyi and Chen, Zhengyang and Chen, Zhuo and Liu, Shujie and Wu, Jian and Qian, Yao and Wei, Furu and Li, Jinyu and others},
+        booktitle={ICASSP 2022-2022 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)},
+        pages={6152--6156},
+        year={2022},
+        organization={IEEE}
+    }
+
+
+unispeech_sat
+~~~~~~~~~~~~~~~~~~~~~
+
+Alias of `unispeech_sat_base`_
+
+
+unispeech_sat_base
+~~~~~~~~~~~~~~~~~~~~~~
+
+- Model Architecture: 12 layers Transformer blocks
+- Unlabled Speech: LibriSpeech 960 hours
+
+
+unispeech_sat_base_plus
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+- Model Architecture: 12 layers Transformer blocks
+- Unlabled Speech: LibriLight 60k hours + Gigaspeech 10k hours + VoxPopuli 24k hours = 94k hours
+
+
+unispeech_sat_large
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+- Model Architecture: 24 layers Transformer blocks
+- Unlabled Speech: LibriLight 60k hours + Gigaspeech 10k hours + VoxPopuli 24k hours = 94k hours
+
+
+
+WavLM
+--------------------------------------------------
+`WavLM: Large-Scale Self-Supervised Pre-Training for Full Stack Speech Processing <https://arxiv.org/abs/2110.13900>`_
+
+.. code-block:: bash
+
+    @article{Chen2021WavLM,
+        title   = {WavLM: Large-Scale Self-Supervised  Pre-training   for Full Stack Speech Processing},
+        author  = {Sanyuan Chen and Chengyi Wang and Zhengyang Chen and Yu Wu and Shujie Liu and Zhuo Chen and Jinyu Li and Naoyuki Kanda and Takuya Yoshioka and Xiong Xiao and Jian Wu and Long Zhou and Shuo Ren and Yanmin Qian and Yao Qian and Jian Wu and Michael Zeng and Furu Wei},
+        eprint={2110.13900},
+        archivePrefix={arXiv},
+        primaryClass={cs.CL},
+        year={2021}
+    }
+
+
+wavlm
+~~~~~~~~~~~~~~~~~
+
+Alias of `wavlm_base_plus`_
+
+
+wavlm_base
+~~~~~~~~~~~~~~~~
+
+- Model Architecture: 12 layers Transformer blocks
+- Unlabled Speech: LibriSpeech 960 hours
+
+
+wavlm_base_plus
+~~~~~~~~~~~~~~~~~~~~~
+
+- Model Architecture: 12 layers Transformer blocks
+- Unlabled Speech: LibriLight 60k hours + Gigaspeech 10k hours + VoxPopuli 24k hours = 94k hours
+
+
+wavlm_large
+~~~~~~~~~~~~~~~~~~~~~
+
+- Model Architecture: 24 layers Transformer blocks
+- Unlabled Speech: LibriLight 60k hours + Gigaspeech 10k hours + VoxPopuli 24k hours = 94k hours
+
+
+data2vec
+--------------------------------------------------
+`data2vec: A General Framework for Self-supervised Learning in Speech, Vision and Language <https://arxiv.org/abs/2202.03555>`_
+
+.. code-block:: bash
+
+    @article{baevski2022data2vec,
+        title={Data2vec: A general framework for self-supervised learning in speech, vision and language},
+        author={Baevski, Alexei and Hsu, Wei-Ning and Xu, Qiantong and Babu, Arun and Gu, Jiatao and Auli, Michael},
+        journal={arXiv preprint arXiv:2202.03555},
+        year={2022}
+    }
+
+
+data2vec
+~~~~~~~~~~~~~~~~~
+
+Alias of `data2vec_base_960`_
+
+
+data2vec_base_960
+~~~~~~~~~~~~~~~~~~
+
+- Model Architecture: 12 layers Transformer blocks
+- Unlabled Speech: LibriSpeech 960 hours
+
+
+data2vec_large_ll60k
+~~~~~~~~~~~~~~~~~~~~~
+
+- Model Architecture: 24 layers Transformer blocks
+- Unlabled Speech: LibriLight 60k hours
+
+
+AST
+--------------------------------------------------
+`AST: Audio Spectrogram Transformer <https://arxiv.org/abs/2104.01778>`_
+
+.. code-block:: bash
+
+    @article{gong2021ast,
+        title={Ast: Audio spectrogram transformer},
+        author={Gong, Yuan and Chung, Yu-An and Glass, James},
+        journal={arXiv preprint arXiv:2104.01778},
+        year={2021}
+    }
+
+
+All the entries below support the following :code:`extra_conf`:
+
+====================  ====================
+column                description
+====================  ====================
+window_secs           (float) -
+                        The segment waveform length to feed into the
+                        AST model. If the input waveform is longer than this
+                        length, do sliding windowing on the waveform and concat
+                        the results along the time axis.
+stride_secs           (float) -
+                        When doing sliding window on the waveform (see
+                        above), the stride seconds between windows.
+====================  ====================
+
+
+ast
+~~~~~~~~~~~~~~~~~~
+
+- Labeled Data: AudioSet
+
+
+SSAST
+--------------------------------------------------
+`SSAST: Self-Supervised Audio Spectrogram Transformer <https://arxiv.org/abs/2110.09784>`_
+
+.. code-block:: bash
+
+    @inproceedings{gong2022ssast,
+        title={Ssast: Self-supervised audio spectrogram transformer},
+        author={Gong, Yuan and Lai, Cheng-I and Chung, Yu-An and Glass, James},
+        booktitle={Proceedings of the AAAI Conference on Artificial Intelligence},
+        volume={36},
+        number={10},
+        pages={10699--10709},
+        year={2022}
+    }
+
+
+All the entries below support the following :code:`extra_conf`:
+
+====================  ====================
+column                description
+====================  ====================
+window_secs           (float) -
+                        The segment waveform length to feed into the
+                        AST model. If the input waveform is longer than this
+                        length, do sliding windowing on the waveform and concat
+                        the results along the time axis.
+====================  ====================
+
+
+ssast_frame_base
+~~~~~~~~~~~~~~~~~~
+
+- Unlabled Data: LibriSpeech & AudioSet
+- fbank patch size: 128 (freq) * 2 (time)
+
+ssast_patch_base
+~~~~~~~~~~~~~~~~~~~
+
+- Unlabled Data: LibriSpeech & AudioSet
+- fbank patch size: 16 (freq) * 16 (time)
+
+
+MAE-AST
+--------------------------------------------------
+`MAE-AST: Masked Autoencoding Audio Spectrogram Transformer <https://arxiv.org/abs/2203.16691>`_
+
+.. code-block:: bash
+
+    @article{baade2022mae,
+        title={MAE-AST: Masked Autoencoding Audio Spectrogram Transformer},
+        author={Baade, Alan and Peng, Puyuan and Harwath, David},
+        journal={arXiv preprint arXiv:2203.16691},
+        year={2022}
+    }
+
+
+mae_ast_frame
+~~~~~~~~~~~~~~~~~~
+
+- Unlabled Data: LibriSpeech & AudioSet
+- fbank patch size: 128 (freq) * 2 (time)
+
+
+mae_ast_patch
+~~~~~~~~~~~~~~~~~~
+
+- Unlabled Data: LibriSpeech & AudioSet
+- fbank patch size: 16 (freq) * 16 (time)
+
+
+Byol-A
+--------------------------------------------------
+`BYOL for Audio: Self-Supervised Learning for General-Purpose Audio Representation <https://arxiv.org/abs/2103.06695>`_
+
+.. code-block:: bash
+
+    @inproceedings{niizumi2021byol,
+        title={BYOL for audio: Self-supervised learning for general-purpose audio representation},
+        author={Niizumi, Daisuke and Takeuchi, Daiki and Ohishi, Yasunori and Harada, Noboru and Kashino, Kunio},
+        booktitle={2021 International Joint Conference on Neural Networks (IJCNN)},
+        pages={1--8},
+        year={2021},
+        organization={IEEE}
+    }
+
+
+All the entries below support the following :code:`extra_conf`:
+
+====================  ====================
+column                description
+====================  ====================
+window_secs           (float) -
+                        The segment waveform length to feed into the
+                        AST model. If the input waveform is longer than this
+                        length, do sliding windowing on the waveform and concat
+                        the results along the time axis.
+stride_secs           (float) -
+                        When doing sliding window on the waveform (see
+                        above), the stride seconds between windows.
+====================  ====================
+
+
+byol_a_2048
+~~~~~~~~~~~~~~~~~~
+
+- Unlabled Data: AudioSet
+
+
+byol_a_1024
+~~~~~~~~~~~~~~~~~~
+
+- Unlabled Data: AudioSet
+
+
+byol_a_512
+~~~~~~~~~~~~~~~~~~
+
+- Unlabled Data: AudioSet
+
+
+Byol-S
+--------------------------------------------------
+`BYOL-S: Learning Self-supervised Speech Representations by Bootstrapping <https://arxiv.org/abs/2206.12038>`_
+
+.. code-block:: bash
+
+    @article{elbanna2022byol,
+        title={Byol-s: Learning self-supervised speech representations by bootstrapping},
+        author={Elbanna, Gasser and Scheidwasser-Clow, Neil and Kegler, Mikolaj and Beckmann, Pierre and Hajal, Karl El and Cernak, Milos},
+        journal={arXiv preprint arXiv:2206.12038},
+        year={2022}
+    }
+
+
+byol_s_default
+~~~~~~~~~~~~~~~~~~
+
+- Unlabled Data: AudioSet (Speech subset)
+
+
+byol_s_cvt
+~~~~~~~~~~~~~~~~~~
+
+- Unlabled Data: AudioSet (Speech subset)
+
+
+byol_s_resnetish34
+~~~~~~~~~~~~~~~~~~
+
+- Unlabled Data: AudioSet (Speech subset)
+
+
+VGGish
+--------------------------------------------------
+`CNN Architectures for Large-Scale Audio Classification <https://arxiv.org/abs/1609.09430>`_
+
+.. code-block:: bash
+
+    @inproceedings{hershey2017cnn,
+        title={CNN architectures for large-scale audio classification},
+        author={Hershey, Shawn and Chaudhuri, Sourish and Ellis, Daniel PW and Gemmeke, Jort F and Jansen, Aren and Moore, R Channing and Plakal, Manoj and Platt, Devin and Saurous, Rif A and Seybold, Bryan and others},
+        booktitle={2017 ieee international conference on acoustics, speech and signal processing (icassp)},
+        pages={131--135},
+        year={2017},
+        organization={IEEE}
+    }
+
+
+vggish
+~~~~~~~~~~~~~~~~~~
+
+- Labaled Data: AudioSet
+
+
+PaSST
+--------------------------------------------------
+`Efficient Training of Audio Transformers with Patchout <https://arxiv.org/abs/2110.05069>`_
+
+.. code-block:: bash
+
+    @article{koutini2021efficient,
+        title={Efficient training of audio transformers with patchout},
+        author={Koutini, Khaled and Schl{\"u}ter, Jan and Eghbal-zadeh, Hamid and Widmer, Gerhard},
+        journal={arXiv preprint arXiv:2110.05069},
+        year={2021}
+    }
+
+All the entries below support the following :code:`extra_conf`:
+
+====================  ====================
+column                description
+====================  ====================
+window_secs           (float) -
+                        The segment waveform length to feed into the
+                        model. If the input waveform is longer than this
+                        length, do sliding windowing on the waveform and concat
+                        the results along the time axis.
+stride_secs           (float) -
+                        When doing sliding window on the waveform (see
+                        above), the stride seconds between windows.
+====================  ====================
+
+passt_base
+~~~~~~~~~~~~~~~~~~
+
+- Labaled Data: AudioSet
+
+
+Authors:
+
+- Leo 2022

s3prl_s3prl_main/docs/util/is_valid.py

@@ -0,0 +1,21 @@
+import os
+import argparse
+from pathlib import Path
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("module_root")
+    parser.add_argument("valid_paths")
+    args = parser.parse_args()
+
+    with open(args.valid_paths) as file:
+        valid_paths = [line.strip() for line in file.readlines()]
+
+    ignored_paths = []
+    module_root_name = Path(args.module_root).stem
+    for item in os.listdir(args.module_root):
+        pattern = f"{module_root_name}/{item}"
+        if pattern not in valid_paths:
+            ignored_paths.append(pattern)
+
+    print(" ".join(ignored_paths))

s3prl_s3prl_main/example/customize.py

@@ -0,0 +1,43 @@
+import pandas as pd
+import torch
+
+from s3prl.problem import SuperbASR
+
+
+class LowResourceLinearSuperbASR(SuperbASR):
+    def prepare_data(
+        self, prepare_data: dict, target_dir: str, cache_dir: str, get_path_only=False
+    ):
+        train_path, valid_path, test_paths = super().prepare_data(
+            prepare_data, target_dir, cache_dir, get_path_only
+        )
+
+        # Take only the first 100 utterances for training
+        df = pd.read_csv(train_path)
+        df = df.iloc[:100]
+        df.to_csv(train_path, index=False)
+
+        return train_path, valid_path, test_paths
+
+    def build_downstream(
+        self,
+        build_downstream: dict,
+        downstream_input_size: int,
+        downstream_output_size: int,
+        downstream_input_stride: int,
+    ):
+        import torch
+
+        class Model(torch.nn.Module):
+            def __init__(self, input_size, output_size) -> None:
+                super().__init__()
+                self.linear = torch.nn.Linear(input_size, output_size)
+
+            def forward(self, x, x_len):
+                return self.linear(x), x_len
+
+        return Model(downstream_input_size, downstream_output_size)
+
+
+if __name__ == "__main__":
+    LowResourceLinearSuperbASR().main()

s3prl_s3prl_main/example/run_asr.sh

@@ -0,0 +1,3 @@
+#!/bin/bash
+
+python3 s3prl/main.py SuperbASR --target_dir result/asr --prepare_data.dataset_root /home/leo/d/datasets/LibriSpeech/ --build_upstream.name apc

s3prl_s3prl_main/example/run_sid.sh

@@ -0,0 +1,4 @@
+#!/bin/bash
+
+python3 s3prl/main.py SuperbSID --target_dir result/tmp/ic --prepare_data.dataset_root /home/leo/d/datasets/VoxCeleb1/ --build_upstream.name apc
+

s3prl_s3prl_main/example/ssl/pretrain.py

@@ -0,0 +1,285 @@
+import argparse
+import logging
+import math
+import os
+from copy import deepcopy
+from pathlib import Path
+
+import torch
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from s3prl import Container, Logs, Object, Output
+from s3prl.dataset.base import AugmentedDynamicItemDataset, DataLoader
+from s3prl.sampler import DistributedBatchSamplerWrapper
+from s3prl.util.configuration import parse_override, qualname_to_cls
+from s3prl.util.seed import fix_random_seeds
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+logger = logging.getLogger(__name__)
+
+DRYRUN_CONFIG = dict(
+    Trainer=dict(
+        total_steps=200000,
+        log_step=5000,
+        valid_step=5000,
+        save_step=5000,
+        eval_batch=8,
+    ),
+)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "problem", help="The problem module. E.g. `s3prl.problem.ssl.tera.Tera`"
+    )
+    parser.add_argument("dataset_root", help="The dataset root for pretrain.")
+    parser.add_argument("save_to", help="The directory to save checkpoint")
+    parser.add_argument("--n_jobs", type=int, default=8)
+    parser.add_argument(
+        "--override",
+        default=None,
+        help=(
+            "Override the default_config of the problem module. "
+            "E.g. --override ValidSampler.batch_size=4,,TestSampler.batch_size=4"
+        ),
+    )
+    parser.add_argument("--resume", action="store_true")
+    parser.add_argument("--dryrun", action="store_true")
+    parser.add_argument("--seed", type=int, default=1337)
+    args = parser.parse_args()
+
+    fix_random_seeds(args.seed)
+    problem = qualname_to_cls(args.problem)
+    config = Container(deepcopy(problem.default_config))
+
+    for key, value in vars(args).items():
+        if key not in ["override"]:
+            config[key] = value
+
+    if args.dryrun:
+        config.override(DRYRUN_CONFIG)
+
+    if isinstance(args.override, str) and len(args.override) > 0:
+        override_dict = parse_override(args.override)
+        config.override(override_dict)
+
+    return problem, config
+
+
+def main():
+    logging.basicConfig(level=logging.INFO)
+
+    problem, config = parse_args()
+    save_to = Path(config.save_to)
+    save_to.mkdir(exist_ok=True, parents=True)
+
+    # configure any upstream
+    body = problem.Body(**config.Body)
+    head = problem.Head(**config.Head)
+    loss = problem.Loss(**config.Loss)
+    stats = Container()
+
+    logger.info("Preparing corpus")
+    corpus = problem.Corpus(config.dataset_root, **config.Corpus)
+    train_data, valid_data, test_data, corpus_stats = corpus().split(3)
+    stats.add(corpus_stats)
+
+    logger.info("Preparing train data")
+    train_dataset = AugmentedDynamicItemDataset(train_data, tools=stats)
+    train_dataset = problem.TrainData(**config.TrainData)(train_dataset)
+    assert train_dataset.get_tool("feat_dim") == problem.input_size
+    train_sampler = DistributedBatchSamplerWrapper(
+        problem.TrainSampler(train_dataset, **config.TrainSampler),
+        num_replicas=1,
+        rank=0,
+    )
+    train_dataloader = DataLoader(
+        train_dataset,
+        train_sampler,
+        num_workers=config.n_jobs,
+    )
+    stats.add(train_dataset.all_tools())
+
+    logger.info("Preparing valid data")
+    valid_dataset = AugmentedDynamicItemDataset(valid_data, tools=stats)
+    valid_dataset = problem.ValidData(**config.ValidData)(valid_dataset)
+    valid_sampler = DistributedBatchSamplerWrapper(
+        problem.ValidSampler(valid_dataset, **config.ValidSampler),
+        num_replicas=1,
+        rank=0,
+    )
+    valid_dataloader = DataLoader(
+        valid_dataset,
+        valid_sampler,
+        num_workers=12,
+    )
+
+    logger.info("Preparing test data")
+    test_dataset = AugmentedDynamicItemDataset(test_data, tools=stats)
+    test_dataset = problem.TestData(**config.TestData)(test_dataset)
+    test_sampler = DistributedBatchSamplerWrapper(
+        problem.ValidSampler(test_dataset, **config.TestSampler),
+        num_replicas=1,
+        rank=0,
+    )
+    test_dataloader = DataLoader(
+        test_dataset,
+        test_sampler,
+        num_workers=12,
+    )
+
+    sorted_ckpt_dirs = sorted(
+        [
+            file
+            for file in save_to.iterdir()
+            if file.is_dir() and str(file).endswith(".ckpts")
+        ],
+        key=os.path.getmtime,
+    )
+
+    if config.resume and len(sorted_ckpt_dirs) > 0:
+        logger.info("Last checkpoint found. Load model and optimizer from checkpoint")
+        task = Object.load_checkpoint(sorted_ckpt_dirs[1] / "task.ckpt").to(device)
+    else:
+        logger.info("Create a new model")
+        task = problem.Task(body, head, loss, **stats)
+        task = task.to(device)
+
+    # ALL THE FOLLOWING CODES ARE FOR TRAINER
+    # WHICH CAN BE LARGELY SIMPLIFIED WHEN USING OTHER TRAINER PACKAGES
+
+    opt_cls_qualname, opt_cfgs = config.Optimizer.split(1)
+    optimizer = qualname_to_cls(opt_cls_qualname)(task.parameters(), **opt_cfgs)
+    if config.resume and len(sorted_ckpt_dirs) > 0:
+        optimizer.load_state_dict(torch.load(sorted_ckpt_dirs[-1] / "optimizer.ckpt"))
+
+    if config.Trainer.use_valid:
+        if config.resume and len(sorted_ckpt_dirs) > 0:
+            valid_best_score = torch.load(
+                sorted_ckpt_dirs[-1] / "valid_best_score.ckpt"
+            )[config.Trainer.valid_metric]
+        else:
+            valid_best_score = -100000 if config.Trainer.valid_higher_better else 100000
+
+    def save_checkpoint(name):
+        ckpt_dir: Path = save_to / f"{name}.ckpts"
+        ckpt_dir.mkdir(parents=True, exist_ok=True)
+        logger.info(f"Save checkpoint to: {ckpt_dir}")
+
+        if hasattr(problem, "save_checkpoint"):
+            logger.info(f"Save upstream checkpoint to: {ckpt_dir}")
+            problem.save_checkpoint(config, body, head, ckpt_dir / "upstream.ckpt")
+        task.save_checkpoint(ckpt_dir / "task.ckpt")
+        torch.save(optimizer.state_dict(), ckpt_dir / "optimizer.ckpt")
+        torch.save(
+            {config.Trainer.valid_metric: valid_best_score},
+            ckpt_dir / "valid_best_score.ckpt",
+        )
+
+    pbar = tqdm(total=config.Trainer.total_steps, desc="Total")
+    train_completed = False
+    accum_grad_steps = 0
+    while not train_completed:
+        batch_results = []
+        for batch in tqdm(train_dataloader, desc="Train", total=len(train_dataloader)):
+            pbar.update(1)
+            global_step = pbar.n
+
+            assert isinstance(batch, Output)
+            batch = batch.to(device)
+
+            task.train()
+            result = task.train_step(**batch)
+            assert isinstance(result, Output)
+
+            result.loss /= config.Trainer.gradient_accumulate_steps
+            result.loss.backward()
+
+            grad_norm = torch.nn.utils.clip_grad_norm_(
+                task.parameters(), max_norm=config.Trainer.gradient_clipping
+            )
+
+            if math.isnan(grad_norm):
+                logger.warning(f"Grad norm is NaN at step {global_step}")
+                optimizer.zero_grad()
+                accum_grad_steps = 0
+            else:
+                accum_grad_steps += 1
+                if accum_grad_steps == config.Trainer.gradient_accumulate_steps:
+                    optimizer.step()
+                    optimizer.zero_grad()
+                    accum_grad_steps = 0
+                batch_results.append(result.cacheable())
+
+            if global_step % config.Trainer.log_step == 0:
+                logs: Logs = task.train_reduction(batch_results).logs
+                logger.info(f"[Train] step {global_step}")
+                for name, value in logs.Scalar.items():
+                    if name == "loss":
+                        value *= config.Trainer.gradient_accumulate_steps
+                    logger.info(f"{name}: {value}")
+                batch_results = []
+
+            if global_step % config.Trainer.valid_step == 0:
+                with torch.no_grad():
+                    if config.Trainer.use_valid:
+                        valid_results = []
+                        for batch_idx, batch in enumerate(
+                            tqdm(
+                                valid_dataloader,
+                                desc="Valid",
+                                total=len(valid_dataloader),
+                            )
+                        ):
+                            if batch_idx == config.Trainer.get("eval_batch", -1):
+                                break
+                            batch = batch.to(device)
+                            task.eval()
+                            result = task.valid_step(**batch)
+                            valid_results.append(result.cacheable())
+
+                        logs: Logs = task.valid_reduction(valid_results).slice(1)
+                        logger.info(f"[Valid] step {global_step}")
+                        for name, value in logs.Scalar.items():
+                            logger.info(f"{name}: {value}")
+                            if name == config.Trainer.valid_metric:
+                                cond1 = config.Trainer.valid_higher_better and (
+                                    value > valid_best_score
+                                )
+                                cond2 = (not config.Trainer.valid_higher_better) and (
+                                    value < valid_best_score
+                                )
+                                if cond1 or cond2:
+                                    valid_best_score = value
+                                    save_checkpoint("valid_best")
+
+            if (
+                global_step % config.Trainer.save_step == 0
+                or global_step == config.Trainer.total_steps
+            ):
+                save_checkpoint(f"global_step_{global_step}")
+
+                if global_step == config.Trainer.total_steps:
+                    train_completed = True
+                    break
+
+    test_results = []
+    for batch_idx, batch in enumerate(
+        tqdm(test_dataloader, desc="Test", total=len(test_dataloader))
+    ):
+        if batch_idx == config.Trainer.get("eval_batch", -1):
+            break
+        batch = batch.to(device)
+        result = task.test_step(**batch)
+        test_results.append(result.cacheable())
+
+    logs: Logs = task.test_reduction(test_results).slice(1)
+    logger.info(f"[Test] step {global_step}")
+    for name, value in logs.Scalar.items():
+        logger.info(f"{name}: {value}")
+
+
+if __name__ == "__main__":
+    main()

s3prl_s3prl_main/example/superb/train.py

@@ -0,0 +1,291 @@
+import argparse
+import logging
+import math
+import os
+from copy import deepcopy
+from pathlib import Path
+
+import torch
+from tqdm import tqdm
+
+from s3prl import Container, Logs, Object, Output
+from s3prl.dataset.base import AugmentedDynamicItemDataset, DataLoader
+from s3prl.nn import S3PRLUpstream, UpstreamDownstreamModel
+from s3prl.sampler import DistributedBatchSamplerWrapper
+from s3prl.util.configuration import parse_override, qualname_to_cls
+from s3prl.util.seed import fix_random_seeds
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+logger = logging.getLogger(__name__)
+
+DRYRUN_CONFIG = dict(
+    Trainer=dict(
+        total_steps=10,
+        log_step=2,
+        valid_step=5,
+        save_step=5,
+        eval_batch=5,
+    ),
+)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("upstream", help="The upstream name. E.g. wav2vec2")
+    parser.add_argument(
+        "problem",
+        help="The problem module. E.g. s3prl.problem.SuperbSID",
+    )
+    parser.add_argument(
+        "dataset_root",
+        help="The dataset root of your problem.",
+    )
+    parser.add_argument("save_to", help="The directory to save checkpoint")
+    parser.add_argument("--feature_selection", default="hidden_states")
+    parser.add_argument("--n_jobs", type=int, default=6)
+    parser.add_argument(
+        "--override",
+        default=None,
+        help=(
+            "Override the default_config of the problem module. "
+            "E.g. --override ValidSampler.batch_size=4,,TestSampler.batch_size=4"
+        ),
+    )
+    parser.add_argument("--resume", action="store_true")
+    parser.add_argument("--dryrun", action="store_true")
+    parser.add_argument("--seed", type=int, default=1337)
+    args = parser.parse_args()
+
+    fix_random_seeds(args.seed)
+    problem = qualname_to_cls(args.problem)
+    config = Container(deepcopy(problem.default_config))
+
+    for key, value in vars(args).items():
+        if key not in ["override"]:
+            config[key] = value
+
+    if args.dryrun:
+        config.override(DRYRUN_CONFIG)
+
+    if isinstance(args.override, str) and len(args.override) > 0:
+        override_dict = parse_override(args.override)
+        config.override(override_dict)
+
+    return problem, config
+
+
+def main():
+    logging.basicConfig(level=logging.INFO)
+
+    problem, config = parse_args()
+    save_to = Path(config.save_to)
+    save_to.mkdir(exist_ok=True, parents=True)
+
+    # configure any upstream
+    upstream = S3PRLUpstream(config.upstream, config.feature_selection)
+    stats = Container(upstream_rate=upstream.downsample_rate)
+
+    logger.info("Preparing corpus")
+    corpus = problem.Corpus(config.dataset_root, **config.Corpus)
+    train_data, valid_data, test_data, corpus_stats = corpus().split(3)
+    stats.add(corpus_stats)
+
+    logger.info("Preparing train data")
+    train_dataset = AugmentedDynamicItemDataset(train_data, tools=stats)
+    train_dataset = problem.TrainData(**config.TrainData)(train_dataset)
+    train_sampler = DistributedBatchSamplerWrapper(
+        problem.TrainSampler(train_dataset, **config.TrainSampler),
+        num_replicas=1,
+        rank=0,
+    )
+    train_dataloader = DataLoader(
+        train_dataset,
+        train_sampler,
+        num_workers=config.n_jobs,
+    )
+    stats.add(train_dataset.all_tools())
+
+    logger.info("Preparing valid data")
+    valid_dataset = AugmentedDynamicItemDataset(valid_data, tools=stats)
+    valid_dataset = problem.ValidData(**config.ValidData)(valid_dataset)
+    valid_sampler = DistributedBatchSamplerWrapper(
+        problem.ValidSampler(valid_dataset, **config.ValidSampler),
+        num_replicas=1,
+        rank=0,
+    )
+    valid_dataloader = DataLoader(
+        valid_dataset,
+        valid_sampler,
+        num_workers=12,
+    )
+
+    logger.info("Preparing test data")
+    test_dataset = AugmentedDynamicItemDataset(test_data, tools=stats)
+    test_dataset = problem.TestData(**config.TestData)(test_dataset)
+    test_sampler = DistributedBatchSamplerWrapper(
+        problem.ValidSampler(test_dataset, **config.TestSampler),
+        num_replicas=1,
+        rank=0,
+    )
+    test_dataloader = DataLoader(
+        test_dataset,
+        test_sampler,
+        num_workers=12,
+    )
+
+    sorted_ckpt_dirs = sorted(
+        [
+            file
+            for file in save_to.iterdir()
+            if file.is_dir() and str(file).endswith(".ckpts")
+        ],
+        key=os.path.getmtime,
+    )
+
+    if config.resume and len(sorted_ckpt_dirs) > 0:
+        logger.info("Last checkpoint found. Load model and optimizer from checkpoint")
+        task = Object.load_checkpoint(sorted_ckpt_dirs[1] / "task.ckpt").to(device)
+    else:
+        logger.info("Create a new model")
+        downstream = problem.Downstream(
+            upstream.output_size,
+            **stats,
+        )
+        model = UpstreamDownstreamModel(upstream, downstream)
+        # task = problem.Task(model, **{**stats, **config.Task})
+        task = problem.Task(model, **stats, **config.Task)
+        task = task.to(device)
+
+    # ALL THE FOLLOWING CODES ARE FOR TRAINER
+    # WHICH CAN BE LARGELY SIMPLIFIED WHEN USING OTHER TRAINER PACKAGES
+
+    opt_cls_qualname, opt_cfgs = config.Optimizer.split(1)
+    optimizer = qualname_to_cls(opt_cls_qualname)(task.parameters(), **opt_cfgs)
+    if config.resume and len(sorted_ckpt_dirs) > 0:
+        optimizer.load_state_dict(torch.load(sorted_ckpt_dirs[-1] / "optimizer.ckpt"))
+
+    if config.Trainer.use_valid:
+        if config.resume and len(sorted_ckpt_dirs) > 0:
+            valid_best_score = torch.load(
+                sorted_ckpt_dirs[-1] / "valid_best_score.ckpt"
+            )[config.Trainer.valid_metric]
+        else:
+            valid_best_score = -100000 if config.Trainer.valid_higher_better else 100000
+
+    def save_checkpoint(name):
+        ckpt_dir: Path = save_to / f"{name}.ckpts"
+        ckpt_dir.mkdir(parents=True, exist_ok=True)
+        logger.info(f"Save checkpoint to: {ckpt_dir}")
+
+        task.save_checkpoint(ckpt_dir / "task.ckpt")
+        torch.save(optimizer.state_dict(), ckpt_dir / "optimizer.ckpt")
+        torch.save(
+            {config.Trainer.valid_metric: valid_best_score},
+            ckpt_dir / "valid_best_score.ckpt",
+        )
+
+    pbar = tqdm(total=config.Trainer.total_steps, desc="Total")
+    train_completed = False
+    accum_grad_steps = 0
+    while not train_completed:
+        batch_results = []
+        for batch in tqdm(train_dataloader, desc="Train", total=len(train_dataloader)):
+            pbar.update(1)
+            global_step = pbar.n
+
+            assert isinstance(batch, Output)
+            batch = batch.to(device)
+
+            task.train()
+            result = task.train_step(**batch)
+            assert isinstance(result, Output)
+
+            result.loss /= config.Trainer.gradient_accumulate_steps
+            result.loss.backward()
+
+            grad_norm = torch.nn.utils.clip_grad_norm_(
+                task.parameters(), max_norm=config.Trainer.gradient_clipping
+            )
+
+            if math.isnan(grad_norm):
+                logger.warning(f"Grad norm is NaN at step {global_step}")
+                optimizer.zero_grad()
+                accum_grad_steps = 0
+            else:
+                accum_grad_steps += 1
+                if accum_grad_steps == config.Trainer.gradient_accumulate_steps:
+                    optimizer.step()
+                    optimizer.zero_grad()
+                    accum_grad_steps = 0
+                batch_results.append(result.cacheable())
+
+            if global_step % config.Trainer.log_step == 0:
+                logs: Logs = task.train_reduction(batch_results).logs
+                logger.info(f"[Train] step {global_step}")
+                for name, value in logs.Scalar.items():
+                    if name == "loss":
+                        value *= config.Trainer.gradient_accumulate_steps
+                    logger.info(f"{name}: {value}")
+                batch_results = []
+
+            if global_step % config.Trainer.valid_step == 0:
+                with torch.no_grad():
+                    if config.Trainer.use_valid:
+                        valid_results = []
+                        for batch_idx, batch in enumerate(
+                            tqdm(
+                                valid_dataloader,
+                                desc="Valid",
+                                total=len(valid_dataloader),
+                            )
+                        ):
+                            if batch_idx == config.Trainer.get("eval_batch", -1):
+                                break
+                            batch = batch.to(device)
+                            task.eval()
+                            result = task.valid_step(**batch)
+                            valid_results.append(result.cacheable())
+
+                        logs: Logs = task.valid_reduction(valid_results).slice(1)
+                        logger.info(f"[Valid] step {global_step}")
+                        for name, value in logs.Scalar.items():
+                            logger.info(f"{name}: {value}")
+                            if name == config.Trainer.valid_metric:
+                                cond1 = config.Trainer.valid_higher_better and (
+                                    value > valid_best_score
+                                )
+                                cond2 = (not config.Trainer.valid_higher_better) and (
+                                    value < valid_best_score
+                                )
+                                if cond1 or cond2:
+                                    valid_best_score = value
+                                    save_checkpoint("valid_best")
+
+            if (
+                global_step % config.Trainer.save_step == 0
+                or global_step == config.Trainer.total_steps
+            ):
+                save_checkpoint(f"global_step_{global_step}")
+
+                if global_step == config.Trainer.total_steps:
+                    train_completed = True
+                    break
+
+    test_results = []
+    for batch_idx, batch in enumerate(
+        tqdm(test_dataloader, desc="Test", total=len(test_dataloader))
+    ):
+        if batch_idx == config.Trainer.get("eval_batch", -1):
+            break
+        batch = batch.to(device)
+        result = task.test_step(**batch)
+        test_results.append(result.cacheable())
+
+    logs: Logs = task.test_reduction(test_results).slice(1)
+    logger.info(f"[Test] step {global_step}")
+    for name, value in logs.Scalar.items():
+        logger.info(f"{name}: {value}")
+
+
+if __name__ == "__main__":
+    main()

s3prl_s3prl_main/example/superb_asr/inference.py

@@ -0,0 +1,40 @@
+import argparse
+from pathlib import Path
+
+import torch
+
+from s3prl import Dataset, Output, Task
+from s3prl.base.object import Object
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "load_from", help="The directory containing all the checkpoints"
+    )
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = parse_args()
+    load_from = Path(args.load_from)
+
+    task: Task = Object.load_checkpoint(load_from / "task.ckpt").to(device)
+    task.eval()
+
+    test_dataset: Dataset = Object.load_checkpoint(load_from / "test_dataset.ckpt")
+    test_dataloader = test_dataset.to_dataloader(batch_size=1, num_workers=6)
+
+    with torch.no_grad():
+        for batch in test_dataloader:
+            batch: Output = batch.to(device)
+            result = task(**batch.subset("x", "x_len", as_type="dict"))
+            for name, prediction in zip(batch.name, result.prediction):
+                print(name, prediction)
+
+
+if __name__ == "__main__":
+    main()

s3prl_s3prl_main/example/superb_asr/train.py

@@ -0,0 +1,241 @@
+import argparse
+import logging
+import math
+from pathlib import Path
+
+import torch
+import torch.optim as optim
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from s3prl import Logs, Object, Output
+from s3prl.nn import S3PRLUpstream, UpstreamDownstreamModel
+from s3prl.sampler import DistributedBatchSamplerWrapper
+from s3prl.superb import asr as problem
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+logger = logging.getLogger(__name__)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("librispeech", help="The root directory of LibriSpeech")
+    parser.add_argument("save_to", help="The directory to save checkpoint")
+    parser.add_argument("--total_steps", type=int, default=200000)
+    parser.add_argument("--log_step", type=int, default=100)
+    parser.add_argument("--eval_step", type=int, default=5000)
+    parser.add_argument("--save_step", type=int, default=100)
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    logging.basicConfig()
+    logger.setLevel(logging.INFO)
+
+    args = parse_args()
+    librispeech = Path(args.librispeech)
+    assert librispeech.is_dir()
+    save_to = Path(args.save_to)
+    save_to.mkdir(exist_ok=True, parents=True)
+
+    logger.info("Preparing preprocessor")
+    preprocessor = problem.Preprocessor(
+        librispeech, splits=["train-clean-100", "dev-clean", "test-clean"]
+    )
+
+    logger.info("Preparing train dataloader")
+    train_dataset = problem.TrainDataset(**preprocessor.train_data())
+    train_sampler = problem.TrainSampler(
+        train_dataset, max_timestamp=16000 * 1000, shuffle=True
+    )
+    train_sampler = DistributedBatchSamplerWrapper(
+        train_sampler, num_replicas=1, rank=0
+    )
+    train_dataloader = DataLoader(
+        train_dataset,
+        batch_sampler=train_sampler,
+        num_workers=4,
+        collate_fn=train_dataset.collate_fn,
+    )
+
+    logger.info("Preparing valid dataloader")
+    valid_dataset = problem.ValidDataset(
+        **preprocessor.valid_data(),
+        **train_dataset.statistics(),
+    )
+    valid_dataset.save_checkpoint(save_to / "valid_dataset.ckpt")
+    valid_sampler = problem.ValidSampler(valid_dataset, 8)
+    valid_sampler = DistributedBatchSamplerWrapper(
+        valid_sampler, num_replicas=1, rank=0
+    )
+    valid_dataloader = DataLoader(
+        valid_dataset,
+        batch_sampler=valid_sampler,
+        num_workers=4,
+        collate_fn=valid_dataset.collate_fn,
+    )
+
+    logger.info("Preparing test dataloader")
+    test_dataset = problem.TestDataset(
+        **preprocessor.test_data(),
+        **train_dataset.statistics(),
+    )
+    test_dataset.save_checkpoint(save_to / "test_dataset.ckpt")
+    test_sampler = problem.TestSampler(test_dataset, 8)
+    test_sampler = DistributedBatchSamplerWrapper(test_sampler, num_replicas=1, rank=0)
+    test_dataloader = DataLoader(
+        test_dataset,
+        batch_sampler=test_sampler,
+        num_workers=4,
+        collate_fn=test_dataset.collate_fn,
+    )
+
+    latest_task = save_to / "task.ckpt"
+    if latest_task.is_file():
+        logger.info("Last checkpoint found. Load model and optimizer from checkpoint")
+
+        # Object.load_checkpoint() from a checkpoint path and
+        # Object.from_checkpoint() from a loaded checkpoint dictionary
+        # are like AutoModel in Huggingface which you only need to
+        # provide the checkpoint for restoring the module.
+        #
+        # Note that source code definition should be importable, since this
+        # auto loading mechanism is just automating the model re-initialization
+        # steps instead of scriptify (torch.jit) all the source code in the
+        # checkpoint
+
+        task = Object.load_checkpoint(latest_task).to(device)
+
+    else:
+        logger.info("No last checkpoint found. Create new model")
+
+        # Model creation block which can be fully customized
+        upstream = S3PRLUpstream("apc")
+        downstream = problem.DownstreamModel(
+            upstream.output_size,
+            preprocessor.statistics().output_size,
+            hidden_size=[512],
+            dropout=[0.2],
+        )
+        model = UpstreamDownstreamModel(upstream, downstream)
+
+        # After customize your own model, simply put it into task object
+        task = problem.Task(model, preprocessor.statistics().label_loader)
+        task = task.to(device)
+
+    # We do not handle optimizer/scheduler in any special way in S3PRL, since
+    # there are lots of dedicated package for this. Hence, we also do not handle
+    # the checkpointing for optimizer/scheduler. Depends on what training pipeline
+    # the user prefer, either Lightning or SpeechBrain, these frameworks will
+    # provide different solutions on how to save these objects. By not handling
+    # these objects in S3PRL we are making S3PRL more flexible and agnostic to training pipeline
+    # The following optimizer codeblock aims to align with the standard usage
+    # of PyTorch which is the standard way to save it.
+
+    optimizer = optim.Adam(task.parameters(), lr=1e-3)
+    latest_optimizer = save_to / "optimizer.ckpt"
+    if latest_optimizer.is_file():
+        optimizer.load_state_dict(torch.load(save_to / "optimizer.ckpt"))
+    else:
+        optimizer = optim.Adam(task.parameters(), lr=1e-3)
+
+    # The following code block demonstrate how to train with your own training loop
+    # This entire block can be easily replaced with Lightning/SpeechBrain Trainer as
+    #
+    #     Trainer(task)
+    #     Trainer.fit(train_dataloader, valid_dataloader, test_dataloader)
+    #
+    # As you can see, there is a huge similarity among train/valid/test loops below,
+    # so it is a natural step to share these logics with a generic Trainer class
+    # as done in Lightning/SpeechBrain
+
+    pbar = tqdm(total=args.total_steps, desc="Total")
+    while True:
+        batch_results = []
+        for batch in tqdm(train_dataloader, desc="Train", total=len(train_dataloader)):
+            pbar.update(1)
+            global_step = pbar.n
+
+            assert isinstance(batch, Output)
+            optimizer.zero_grad()
+
+            # An Output object can more all its direct
+            # attributes/values to the device
+            batch = batch.to(device)
+
+            # An Output object is an OrderedDict so we
+            # can use dict decomposition here
+            task.train()
+            result = task.train_step(**batch)
+            assert isinstance(result, Output)
+
+            # The output of train step must contain
+            # at least a loss key
+            result.loss.backward()
+
+            # gradient clipping
+            grad_norm = torch.nn.utils.clip_grad_norm_(task.parameters(), max_norm=1.0)
+
+            if math.isnan(grad_norm):
+                logger.warning(f"Grad norm is NaN at step {global_step}")
+            else:
+                optimizer.step()
+
+            # Detach from GPU, remove large logging (to Tensorboard or local files)
+            # objects like logits and leave only small data like loss scalars / prediction
+            # strings, so that these objects can be safely cached in a list in the MEM,
+            # and become useful for calculating metrics later
+            # The Output class can do these with self.cacheable()
+            cacheable_result = result.cacheable()
+
+            # Cache these small data for later metric calculation
+            batch_results.append(cacheable_result)
+
+            if (global_step + 1) % args.log_step == 0:
+                logs: Logs = task.train_reduction(batch_results).logs
+                logger.info(f"[Train] step {global_step}")
+                for log in logs.values():
+                    logger.info(f"{log.name}: {log.data}")
+                batch_results = []
+
+            if (global_step + 1) % args.eval_step == 0:
+                with torch.no_grad():
+                    task.eval()
+
+                    # valid
+                    valid_results = []
+                    for batch in tqdm(
+                        valid_dataloader, desc="Valid", total=len(valid_dataloader)
+                    ):
+                        batch = batch.to(device)
+                        result = task.valid_step(**batch)
+                        cacheable_result = result.cacheable()
+                        valid_results.append(cacheable_result)
+
+                    logs: Logs = task.valid_reduction(valid_results).logs
+                    logger.info(f"[Valid] step {global_step}")
+                    for log in logs.values():
+                        logger.info(f"{log.name}: {log.data}")
+
+            if (global_step + 1) % args.save_step == 0:
+                task.save_checkpoint(save_to / "task.ckpt")
+                torch.save(optimizer.state_dict(), save_to / "optimizer.ckpt")
+
+    with torch.no_grad():
+        # test
+        test_results = []
+        for batch in tqdm(test_dataloader, desc="Test", total=len(test_dataloader)):
+            batch = batch.to(device)
+            result = task.test_step(**batch)
+            cacheable_result = result.cacheable()
+            test_results.append(cacheable_result)
+
+        logs: Logs = task.test_reduction(test_results).logs
+        logger.info(f"[Test] step results")
+        for log in logs.values():
+            logger.info(f"{log.name}: {log.data}")
+
+
+if __name__ == "__main__":
+    main()

s3prl_s3prl_main/example/superb_asr/train_with_lightning.py

@@ -0,0 +1,127 @@
+import argparse
+import logging
+from pathlib import Path
+
+import torch
+import torch.optim as optim
+from pytorch_lightning import Trainer
+from pytorch_lightning.callbacks import ModelCheckpoint
+
+from s3prl.nn import S3PRLUpstream, UpstreamDownstreamModel
+from s3prl.superb import asr as problem
+from s3prl.wrapper import LightningModuleSimpleWrapper
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+logger = logging.getLogger(__name__)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("librispeech", help="The root directory of LibriSpeech")
+    parser.add_argument("save_to", help="The directory to save checkpoint")
+    parser.add_argument("--total_steps", type=int, default=200000)
+    parser.add_argument("--log_step", type=int, default=100)
+    parser.add_argument("--eval_step", type=int, default=5000)
+    parser.add_argument("--save_step", type=int, default=100)
+    parser.add_argument(
+        "--not_resume",
+        action="store_true",
+        help="Don't resume from the last checkpoint",
+    )
+
+    # for debugging
+    parser.add_argument("--limit_train_batches", type=int)
+    parser.add_argument("--limit_val_batches", type=int)
+    parser.add_argument("--fast_dev_run", action="store_true")
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    logging.basicConfig(level=logging.INFO)
+
+    args = parse_args()
+    librispeech = Path(args.librispeech)
+    save_to = Path(args.save_to)
+    save_to.mkdir(exist_ok=True, parents=True)
+
+    logger.info("Preparing preprocessor")
+    preprocessor = problem.Preprocessor(librispeech)
+
+    logger.info("Preparing train dataloader")
+    train_dataset = problem.TrainDataset(**preprocessor.train_data())
+    train_dataloader = train_dataset.to_dataloader(
+        batch_size=8,
+        num_workers=6,
+        shuffle=True,
+    )
+
+    logger.info("Preparing valid dataloader")
+    valid_dataset = problem.ValidDataset(
+        **preprocessor.valid_data(),
+        **train_dataset.statistics(),
+    )
+    valid_dataloader = valid_dataset.to_dataloader(batch_size=8, num_workers=6)
+
+    logger.info("Preparing test dataloader")
+    test_dataset = problem.TestDataset(
+        **preprocessor.test_data(),
+        **train_dataset.statistics(),
+    )
+    test_dataloader = test_dataset.to_dataloader(batch_size=8, num_workers=6)
+
+    valid_dataset.save_checkpoint(save_to / "valid_dataset.ckpt")
+    test_dataset.save_checkpoint(save_to / "test_dataset.ckpt")
+
+    upstream = S3PRLUpstream("apc")
+    downstream = problem.DownstreamModel(
+        upstream.output_size, preprocessor.statistics().output_size
+    )
+    model = UpstreamDownstreamModel(upstream, downstream)
+    task = problem.Task(model, preprocessor.statistics().label_loader)
+
+    optimizer = optim.Adam(task.parameters(), lr=1e-3)
+    lightning_task = LightningModuleSimpleWrapper(task, optimizer)
+
+    # The above is the usage of our library
+
+    # The below is pytorch-lightning specific usage, which can be very simple
+    # or very sophisticated, depending on how much you want to customized your
+    # training loop
+
+    checkpoint_callback = ModelCheckpoint(
+        dirpath=str(save_to),
+        filename="superb-asr-{step:02d}-{valid_0_wer:.2f}",
+        monitor="valid_0_wer",  # since might have multiple valid dataloaders
+        save_last=True,
+        save_top_k=3,  # top 3 best ckpt on valid
+        mode="min",  # lower, better
+        every_n_train_steps=args.save_step,
+    )
+
+    trainer = Trainer(
+        callbacks=[checkpoint_callback],
+        accelerator="gpu",
+        gpus=1,
+        max_steps=args.total_steps,
+        log_every_n_steps=args.log_step,
+        val_check_interval=args.eval_step,
+        limit_val_batches=args.limit_val_batches or 1.0,
+        limit_train_batches=args.limit_train_batches or 1.0,
+        fast_dev_run=args.fast_dev_run,
+    )
+
+    last_ckpt = save_to / "last.ckpt"
+    if args.not_resume or not last_ckpt.is_file():
+        last_ckpt = None
+
+    trainer.fit(
+        lightning_task,
+        train_dataloader,
+        val_dataloaders=[valid_dataloader, test_dataloader],
+        ckpt_path=last_ckpt,
+    )
+
+
+if __name__ == "__main__":
+    main()

s3prl_s3prl_main/example/superb_sid/inference.py

@@ -0,0 +1,40 @@
+import argparse
+from pathlib import Path
+
+import torch
+
+from s3prl import Dataset, Output, Task
+from s3prl.base.object import Object
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "load_from", help="The directory containing all the checkpoints"
+    )
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = parse_args()
+    load_from = Path(args.load_from)
+
+    task: Task = Object.load_checkpoint(load_from / "task.ckpt").to(device)
+    task.eval()
+
+    test_dataset: Dataset = Object.load_checkpoint(load_from / "test_dataset.ckpt")
+    test_dataloader = test_dataset.to_dataloader(batch_size=1, num_workers=6)
+
+    with torch.no_grad():
+        for batch in test_dataloader:
+            batch: Output = batch.to(device)
+            result = task(**batch.subset("x", "x_len", as_type="dict"))
+            for name, prediction in zip(batch.name, result.prediction):
+                print(name, prediction)
+
+
+if __name__ == "__main__":
+    main()

s3prl_s3prl_main/example/superb_sid/train.py

@@ -0,0 +1,235 @@
+import argparse
+import logging
+import math
+from pathlib import Path
+
+import torch
+import torch.optim as optim
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from s3prl import Logs, Object, Output
+from s3prl.nn import S3PRLUpstream, UpstreamDownstreamModel
+from s3prl.sampler import DistributedBatchSamplerWrapper
+from s3prl.superb import sid as problem
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+logger = logging.getLogger(__name__)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("voxceleb1", help="The root directory of VoxCeleb1")
+    parser.add_argument("save_to", help="The directory to save checkpoint")
+    parser.add_argument("--total_steps", type=int, default=200000)
+    parser.add_argument("--log_step", type=int, default=100)
+    parser.add_argument("--eval_step", type=int, default=5000)
+    parser.add_argument("--save_step", type=int, default=100)
+    parser.add_argument("--resume", action="store_true")
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    logging.basicConfig()
+    logger.setLevel(logging.INFO)
+
+    args = parse_args()
+    voxceleb1 = Path(args.voxceleb1)
+    assert voxceleb1.is_dir()
+    save_to = Path(args.save_to)
+    save_to.mkdir(exist_ok=True, parents=True)
+
+    logger.info("Preparing preprocessor")
+    preprocessor = problem.Preprocessor(voxceleb1)
+
+    logger.info("Preparing train dataloader")
+    train_dataset = problem.TrainDataset(**preprocessor.train_data())
+    train_sampler = problem.TrainSampler(
+        train_dataset, max_timestamp=16000 * 200, shuffle=True
+    )
+    train_sampler = DistributedBatchSamplerWrapper(
+        train_sampler, num_replicas=1, rank=0
+    )
+    train_dataloader = DataLoader(
+        train_dataset,
+        batch_sampler=train_sampler,
+        num_workers=12,
+        collate_fn=train_dataset.collate_fn,
+    )
+
+    logger.info("Preparing valid dataloader")
+    valid_dataset = problem.ValidDataset(
+        **preprocessor.valid_data(),
+        **train_dataset.statistics(),
+    )
+    valid_dataset.save_checkpoint(save_to / "valid_dataset.ckpt")
+    valid_sampler = problem.ValidSampler(valid_dataset, 8)
+    valid_sampler = DistributedBatchSamplerWrapper(
+        valid_sampler, num_replicas=1, rank=0
+    )
+    valid_dataloader = DataLoader(
+        valid_dataset,
+        batch_sampler=valid_sampler,
+        num_workers=12,
+        collate_fn=valid_dataset.collate_fn,
+    )
+
+    logger.info("Preparing test dataloader")
+    test_dataset = problem.TestDataset(
+        **preprocessor.test_data(),
+        **train_dataset.statistics(),
+    )
+    test_dataset.save_checkpoint(save_to / "test_dataset.ckpt")
+    test_sampler = problem.TestSampler(test_dataset, 8)
+    test_sampler = DistributedBatchSamplerWrapper(test_sampler, num_replicas=1, rank=0)
+    test_dataloader = DataLoader(
+        test_dataset, batch_size=8, num_workers=12, collate_fn=test_dataset.collate_fn
+    )
+
+    latest_task = save_to / "task.ckpt"
+    if args.resume and latest_task.is_file():
+        logger.info("Last checkpoint found. Load model and optimizer from checkpoint")
+
+        # Object.load_checkpoint() from a checkpoint path and
+        # Object.from_checkpoint() from a loaded checkpoint dictionary
+        # are like AutoModel in Huggingface which you only need to
+        # provide the checkpoint for restoring the module.
+        #
+        # Note that source code definition should be importable, since this
+        # auto loading mechanism is just automating the model re-initialization
+        # steps instead of scriptify (torch.jit) all the source code in the
+        # checkpoint
+
+        task = Object.load_checkpoint(latest_task).to(device)
+
+    else:
+        logger.info("No last checkpoint found. Create new model")
+
+        # Model creation block which can be fully customized
+        upstream = S3PRLUpstream("wav2vec2")
+        downstream = problem.DownstreamModel(
+            upstream.output_size, len(preprocessor.statistics().category)
+        )
+        model = UpstreamDownstreamModel(upstream, downstream)
+
+        # After customize your own model, simply put it into task object
+        task = problem.Task(model, preprocessor.statistics().category)
+        task = task.to(device)
+
+    # We do not handle optimizer/scheduler in any special way in S3PRL, since
+    # there are lots of dedicated package for this. Hence, we also do not handle
+    # the checkpointing for optimizer/scheduler. Depends on what training pipeline
+    # the user prefer, either Lightning or SpeechBrain, these frameworks will
+    # provide different solutions on how to save these objects. By not handling
+    # these objects in S3PRL we are making S3PRL more flexible and agnostic to training pipeline
+    # The following optimizer codeblock aims to align with the standard usage
+    # of PyTorch which is the standard way to save it.
+
+    optimizer = optim.Adam(task.parameters(), lr=1e-3)
+    latest_optimizer = save_to / "optimizer.ckpt"
+    if args.resume and latest_optimizer.is_file():
+        optimizer.load_state_dict(torch.load(save_to / "optimizer.ckpt"))
+    else:
+        optimizer = optim.Adam(task.parameters(), lr=1e-3)
+
+    # The following code block demonstrate how to train with your own training loop
+    # This entire block can be easily replaced with Lightning/SpeechBrain Trainer as
+    #
+    #     Trainer(task)
+    #     Trainer.fit(train_dataloader, valid_dataloader, test_dataloader)
+    #
+    # As you can see, there is a huge similarity among train/valid/test loops below,
+    # so it is a natural step to share these logics with a generic Trainer class
+    # as done in Lightning/SpeechBrain
+
+    pbar = tqdm(total=args.total_steps, desc="Total")
+    while True:
+        batch_results = []
+        for batch in tqdm(train_dataloader, desc="Train", total=len(train_dataloader)):
+            pbar.update(1)
+            global_step = pbar.n
+
+            assert isinstance(batch, Output)
+            optimizer.zero_grad()
+
+            # An Output object can more all its direct
+            # attributes/values to the device
+            batch = batch.to(device)
+
+            # An Output object is an OrderedDict so we
+            # can use dict decomposition here
+            task.train()
+            result = task.train_step(**batch)
+            assert isinstance(result, Output)
+
+            # The output of train step must contain
+            # at least a loss key
+            result.loss.backward()
+
+            # gradient clipping
+            grad_norm = torch.nn.utils.clip_grad_norm_(task.parameters(), max_norm=1.0)
+
+            if math.isnan(grad_norm):
+                logger.warning(f"Grad norm is NaN at step {global_step}")
+            else:
+                optimizer.step()
+
+            # Detach from GPU, remove large logging (to Tensorboard or local files)
+            # objects like logits and leave only small data like loss scalars / prediction
+            # strings, so that these objects can be safely cached in a list in the MEM,
+            # and become useful for calculating metrics later
+            # The Output class can do these with self.cacheable()
+            cacheable_result = result.cacheable()
+
+            # Cache these small data for later metric calculation
+            batch_results.append(cacheable_result)
+
+            if (global_step + 1) % args.log_step == 0:
+                logs: Logs = task.train_reduction(batch_results).logs
+                logger.info(f"[Train] step {global_step}")
+                for log in logs.values():
+                    logger.info(f"{log.name}: {log.data}")
+                batch_results = []
+
+            if (global_step + 1) % args.eval_step == 0:
+                with torch.no_grad():
+                    task.eval()
+
+                    # valid
+                    valid_results = []
+                    for batch in tqdm(
+                        valid_dataloader, desc="Valid", total=len(valid_dataloader)
+                    ):
+                        batch = batch.to(device)
+                        result = task.valid_step(**batch)
+                        cacheable_result = result.cacheable()
+                        valid_results.append(cacheable_result)
+
+                    logs: Logs = task.valid_reduction(valid_results).logs
+                    logger.info(f"[Valid] step {global_step}")
+                    for log in logs.values():
+                        logger.info(f"{log.name}: {log.data}")
+
+                    # test
+                    test_results = []
+                    for batch in tqdm(
+                        test_dataloader, desc="Test", total=len(test_dataloader)
+                    ):
+                        batch = batch.to(device)
+                        result = task.test_step(**batch)
+                        cacheable_result = result.cacheable()
+                        test_results.append(cacheable_result)
+
+                    logs: Logs = task.test_reduction(test_results).logs
+                    logger.info(f"[Test] step {global_step}")
+                    for log in logs.values():
+                        logger.info(f"{log.name}: {log.data}")
+
+            if (global_step + 1) % args.save_step == 0:
+                task.save_checkpoint(save_to / "task.ckpt")
+                torch.save(optimizer.state_dict(), save_to / "optimizer.ckpt")
+
+
+if __name__ == "__main__":
+    main()

s3prl_s3prl_main/example/superb_sid/train_with_lightning.py

@@ -0,0 +1,127 @@
+import argparse
+import logging
+from pathlib import Path
+
+import torch
+import torch.optim as optim
+from pytorch_lightning import Trainer
+from pytorch_lightning.callbacks import ModelCheckpoint
+
+from s3prl.nn import S3PRLUpstream, UpstreamDownstreamModel
+from s3prl.superb import sid as problem
+from s3prl.wrapper import LightningModuleSimpleWrapper
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+logger = logging.getLogger(__name__)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("voxceleb1", help="The root directory of VoxCeleb1")
+    parser.add_argument("save_to", help="The directory to save checkpoint")
+    parser.add_argument("--total_steps", type=int, default=200000)
+    parser.add_argument("--log_step", type=int, default=100)
+    parser.add_argument("--eval_step", type=int, default=5000)
+    parser.add_argument("--save_step", type=int, default=100)
+    parser.add_argument(
+        "--not_resume",
+        action="store_true",
+        help="Don't resume from the last checkpoint",
+    )
+
+    # for debugging
+    parser.add_argument("--limit_train_batches", type=int)
+    parser.add_argument("--limit_val_batches", type=int)
+    parser.add_argument("--fast_dev_run", action="store_true")
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    logging.basicConfig(level=logging.INFO)
+
+    args = parse_args()
+    voxceleb1 = Path(args.voxceleb1)
+    save_to = Path(args.save_to)
+    save_to.mkdir(exist_ok=True, parents=True)
+
+    logger.info("Preparing preprocessor")
+    preprocessor = problem.Preprocessor(voxceleb1)
+
+    logger.info("Preparing train dataloader")
+    train_dataset = problem.TrainDataset(**preprocessor.train_data())
+    train_dataloader = train_dataset.to_dataloader(
+        batch_size=8,
+        num_workers=6,
+        shuffle=True,
+    )
+
+    logger.info("Preparing valid dataloader")
+    valid_dataset = problem.ValidDataset(
+        **preprocessor.valid_data(),
+        **train_dataset.statistics(),
+    )
+    valid_dataloader = valid_dataset.to_dataloader(batch_size=8, num_workers=6)
+
+    logger.info("Preparing test dataloader")
+    test_dataset = problem.TestDataset(
+        **preprocessor.test_data(),
+        **train_dataset.statistics(),
+    )
+    test_dataloader = test_dataset.to_dataloader(batch_size=8, num_workers=6)
+
+    valid_dataset.save_checkpoint(save_to / "valid_dataset.ckpt")
+    test_dataset.save_checkpoint(save_to / "test_dataset.ckpt")
+
+    upstream = S3PRLUpstream("apc")
+    downstream = problem.DownstreamModel(
+        upstream.output_size, len(preprocessor.statistics().category)
+    )
+    model = UpstreamDownstreamModel(upstream, downstream)
+    task = problem.Task(model, preprocessor.statistics().category)
+
+    optimizer = optim.Adam(task.parameters(), lr=1e-3)
+    lightning_task = LightningModuleSimpleWrapper(task, optimizer)
+
+    # The above is the usage of our library
+
+    # The below is pytorch-lightning specific usage, which can be very simple
+    # or very sophisticated, depending on how much you want to customized your
+    # training loop
+
+    checkpoint_callback = ModelCheckpoint(
+        dirpath=str(save_to),
+        filename="superb-sid-{step:02d}-{valid_0_accuracy:.2f}",
+        monitor="valid_0_accuracy",  # since might have multiple valid dataloaders
+        save_last=True,
+        save_top_k=3,  # top 3 best ckpt on valid
+        mode="max",  # higher, better
+        every_n_train_steps=args.save_step,
+    )
+
+    trainer = Trainer(
+        callbacks=[checkpoint_callback],
+        accelerator="gpu",
+        gpus=1,
+        max_steps=args.total_steps,
+        log_every_n_steps=args.log_step,
+        val_check_interval=args.eval_step,
+        limit_val_batches=args.limit_val_batches or 1.0,
+        limit_train_batches=args.limit_train_batches or 1.0,
+        fast_dev_run=args.fast_dev_run,
+    )
+
+    last_ckpt = save_to / "last.ckpt"
+    if args.not_resume or not last_ckpt.is_file():
+        last_ckpt = None
+
+    trainer.fit(
+        lightning_task,
+        train_dataloader,
+        val_dataloaders=[valid_dataloader, test_dataloader],
+        ckpt_path=last_ckpt,
+    )
+
+
+if __name__ == "__main__":
+    main()

s3prl_s3prl_main/example/superb_sv/inference.py

@@ -0,0 +1,47 @@
+import argparse
+from pathlib import Path
+
+import torch
+from torch.utils.data import DataLoader
+
+from s3prl import Output
+from s3prl.base.object import Object
+from s3prl.dataset import Dataset
+from s3prl.task import Task
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--load_from",
+        type=str,
+        default="result/sv",
+        help="The directory containing all the checkpoints",
+    )
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = parse_args()
+    load_from = Path(args.load_from)
+
+    task: Task = Object.load_checkpoint(load_from / "task.ckpt").to(device)
+    task.eval()
+
+    test_dataset: Dataset = Object.load_checkpoint(load_from / "test_dataset.ckpt")
+    test_dataloader = DataLoader(
+        test_dataset, batch_size=1, num_workers=6, collate_fn=test_dataset.collate_fn
+    )
+
+    with torch.no_grad():
+        for batch in test_dataloader:
+            batch: Output = batch.to(device)
+            result = task(**batch.subset("x", "x_len", as_type="dict"))
+            print(result.hidden_states.shape)
+
+
+if __name__ == "__main__":
+    main()

s3prl_s3prl_main/example/superb_sv/train.py

@@ -0,0 +1,266 @@
+import argparse
+import logging
+import math
+from pathlib import Path
+
+import torch
+import torch.optim as optim
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from s3prl import Logs, Object, Output
+from s3prl.nn import S3PRLUpstream, UpstreamDownstreamModel
+from s3prl.sampler import DistributedBatchSamplerWrapper
+from s3prl.superb import sv as problem
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+logger = logging.getLogger(__name__)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--voxceleb1",
+        type=str,
+        default="/work/jason410/PublicData/Voxceleb1",
+        help="The root directory of VoxCeleb1",
+    )
+    parser.add_argument(
+        "--save_to",
+        type=str,
+        default="result/sv",
+        help="The directory to save checkpoint",
+    )
+    parser.add_argument("--total_steps", type=int, default=200000)
+    parser.add_argument("--log_step", type=int, default=100)
+    parser.add_argument("--eval_step", type=int, default=200)
+    parser.add_argument("--save_step", type=int, default=100)
+
+    parser.add_argument("--backbone", type=str, default="XVector")
+    parser.add_argument("--pooling_type", type=str, default="TAP")
+    parser.add_argument("--loss_type", type=str, default="softmax")
+    parser.add_argument("--spk_embd_dim", type=int, default=1500)
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    logging.basicConfig()
+    logger.setLevel(logging.INFO)
+
+    args = parse_args()
+    voxceleb1 = Path(args.voxceleb1)
+    assert voxceleb1.is_dir()
+    save_to = Path(args.save_to)
+    save_to.mkdir(exist_ok=True, parents=True)
+
+    logger.info("Preparing preprocessor")
+    preprocessor = problem.Preprocessor(voxceleb1)
+
+    logger.info("Preparing train dataloader")
+    train_dataset = problem.TrainDataset(**preprocessor.train_data())
+    train_sampler = problem.TrainSampler(
+        train_dataset, max_timestamp=16000 * 1000, shuffle=True
+    )
+    train_sampler = DistributedBatchSamplerWrapper(
+        train_sampler, num_replicas=1, rank=0
+    )
+    train_dataloader = DataLoader(
+        train_dataset,
+        batch_sampler=train_sampler,
+        num_workers=6,
+        collate_fn=train_dataset.collate_fn,
+    )
+
+    logger.info("Preparing valid dataloader")
+    valid_dataset = problem.ValidDataset(
+        **preprocessor.valid_data(),
+        **train_dataset.statistics(),
+    )
+    valid_dataset.save_checkpoint(save_to / "valid_dataset.ckpt")
+    # valid_dataset_reload = Object.load_checkpoint(save_to / "valid_dataset.ckpt")
+    valid_sampler = problem.TrainSampler(
+        valid_dataset, max_timestamp=16000 * 1000, shuffle=True
+    )
+    valid_sampler = DistributedBatchSamplerWrapper(
+        valid_sampler, num_replicas=1, rank=0
+    )
+    valid_dataloader = DataLoader(
+        valid_dataset,
+        batch_sampler=valid_sampler,
+        num_workers=6,
+        collate_fn=valid_dataset.collate_fn,
+    )
+
+    logger.info("Preparing test dataloader")
+    test_dataset = problem.TestDataset(
+        **preprocessor.test_data(),
+        **train_dataset.statistics(),
+    )
+    test_dataset.save_checkpoint(save_to / "test_dataset.ckpt")
+    test_sampler = problem.TestSampler(test_dataset, 8)
+    test_sampler = DistributedBatchSamplerWrapper(test_sampler, num_replicas=1, rank=0)
+    test_dataloader = DataLoader(
+        test_dataset, batch_size=1, num_workers=6, collate_fn=test_dataset.collate_fn
+    )
+
+    latest_task = save_to / "task.ckpt"
+    if latest_task.is_file():
+        logger.info("Last checkpoint found. Load model and optimizer from checkpoint")
+
+        # Object.load_checkpoint() from a checkpoint path and
+        # Object.from_checkpoint() from a loaded checkpoint dictionary
+        # are like AutoModel in Huggingface which you only need to
+        # provide the checkpoint for restoring the module.
+        #
+        # Note that source code definition should be importable, since this
+        # auto loading mechanism is just automating the model re-initialization
+        # steps instead of scriptify (torch.jit) all the source code in the
+        # checkpoint
+
+        task = Object.load_checkpoint(latest_task).to(device)
+
+    else:
+        logger.info("No last checkpoint found. Create new model")
+
+        # Model creation block which can be fully customized
+        upstream = S3PRLUpstream("apc")
+        # Have to specify the backbone, pooling_type
+        downstream = problem.speaker_embedding_extractor(
+            backbone=args.backbone,
+            pooling_type=args.pooling_type,
+            input_size=upstream.output_size,
+            output_size=args.spk_embd_dim,
+        )
+
+        model = UpstreamDownstreamModel(upstream, downstream)
+
+        # Have to specify the loss_type
+        task = problem.Task(
+            model=model,
+            categories=preprocessor.statistics().category,
+            loss_type=args.loss_type,
+            trials=test_dataset.statistics().label,
+        )
+        task = task.to(device)
+
+    # We do not handle optimizer/scheduler in any special way in S3PRL, since
+    # there are lots of dedicated package for this. Hence, we also do not handle
+    # the checkpointing for optimizer/scheduler. Depends on what training pipeline
+    # the user prefer, either Lightning or SpeechBrain, these frameworks will
+    # provide different solutions on how to save these objects. By not handling
+    # these objects in S3PRL we are making S3PRL more flexible and agnostic to training pipeline
+    # The following optimizer codeblock aims to align with the standard usage
+    # of PyTorch which is the standard way to save it.
+
+    optimizer = optim.Adam(task.parameters(), lr=1e-3)
+    latest_optimizer = save_to / "optimizer.ckpt"
+
+    if latest_optimizer.is_file():
+        optimizer.load_state_dict(torch.load(save_to / "optimizer.ckpt"))
+    else:
+        optimizer = optim.Adam(task.parameters(), lr=1e-3)
+
+    # The following code block demonstrate how to train with your own training loop
+    # This entire block can be easily replaced with Lightning/SpeechBrain Trainer as
+    #
+    #     Trainer(task)
+    #     Trainer.fit(train_dataloader, valid_dataloader, test_dataloader)
+    #
+    # As you can see, there is a huge similarity among train/valid/test loops below,
+    # so it is a natural step to share these logics with a generic Trainer class
+    # as done in Lightning/SpeechBrain
+
+    pbar = tqdm(total=args.total_steps, desc="Total")
+    while True:
+        batch_results = []
+        for batch in tqdm(train_dataloader, desc="Train", total=len(train_dataloader)):
+            pbar.update(1)
+            global_step = pbar.n
+
+            assert isinstance(batch, Output)
+            optimizer.zero_grad()
+
+            # An Output object can more all its direct
+            # attributes/values to the device
+            batch = batch.to(device)
+
+            # An Output object is an OrderedDict so we
+            # can use dict decomposition here
+            task.train()
+            result = task.train_step(**batch)
+            assert isinstance(result, Output)
+
+            # The output of train step must contain
+            # at least a loss key
+            result.loss.backward()
+
+            # gradient clipping
+            grad_norm = torch.nn.utils.clip_grad_norm_(task.parameters(), max_norm=1.0)
+
+            if math.isnan(grad_norm):
+                logger.warning(f"Grad norm is NaN at step {global_step}")
+            else:
+                optimizer.step()
+
+            # Detach from GPU, remove large logging (to Tensorboard or local files)
+            # objects like logits and leave only small data like loss scalars / prediction
+            # strings, so that these objects can be safely cached in a list in the MEM,
+            # and become useful for calculating metrics later
+            # The Output class can do these with self.cacheable()
+            cacheable_result = result.cacheable()
+
+            # Cache these small data for later metric calculation
+            batch_results.append(cacheable_result)
+
+            if (global_step + 1) % args.log_step == 0:
+                logs: Logs = task.train_reduction(batch_results).logs
+                logger.info(f"[Train] step {global_step}")
+                for log in logs.values():
+                    logger.info(f"{log.name}: {log.data}")
+                batch_results = []
+
+            if (global_step + 1) % args.eval_step == 0:
+                with torch.no_grad():
+                    task.eval()
+
+                    # valid
+                    valid_results = []
+                    for batch in tqdm(
+                        valid_dataloader, desc="Valid", total=len(valid_dataloader)
+                    ):
+                        batch = batch.to(device)
+                        result = task.valid_step(**batch)
+                        cacheable_result = result.cacheable()
+                        valid_results.append(cacheable_result)
+
+                    logs: Logs = task.valid_reduction(valid_results).logs
+                    logger.info(f"[Valid] step {global_step}")
+                    for log in logs.values():
+                        logger.info(f"{log.name}: {log.data}")
+
+                    # test
+                    test_results = []
+                    for batch in tqdm(
+                        test_dataloader, desc="Test", total=len(test_dataloader)
+                    ):
+                        batch = batch.to(device)
+                        result = task.test_step(**batch)
+                        test_results.append(result)
+                        # for key, value in zip(result.name, result.output):
+                        #     test_results[key] = value
+
+                    logs: Logs = task.test_reduction(
+                        batch_results=test_results,
+                    ).logs
+                    logger.info(f"[Test] step {global_step}")
+                    for log in logs.values():
+                        logger.info(f"{log.name}: {log.data}")
+
+            if (global_step + 1) % args.save_step == 0:
+                task.save_checkpoint(save_to / "task.ckpt")
+                torch.save(optimizer.state_dict(), save_to / "optimizer.ckpt")
+
+
+if __name__ == "__main__":
+    main()

s3prl_s3prl_main/example/superb_sv/train_with_lightning.py

@@ -0,0 +1,184 @@
+import argparse
+import logging
+from pathlib import Path
+
+import torch
+import torch.optim as optim
+from pytorch_lightning import Trainer
+from pytorch_lightning.callbacks import ModelCheckpoint
+from torch.utils.data import DataLoader
+
+from s3prl.nn import S3PRLUpstream, UpstreamDownstreamModel
+from s3prl.sampler import DistributedBatchSamplerWrapper
+from s3prl.superb import sv as problem
+from s3prl.wrapper import LightningModuleSimpleWrapper
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+logger = logging.getLogger(__name__)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--voxceleb1",
+        type=str,
+        default="/work/jason410/PublicData/Voxceleb1",
+        help="The root directory of VoxCeleb1",
+    )
+    parser.add_argument(
+        "--save_to",
+        type=str,
+        default="lightning_result/sv",
+        help="The directory to save checkpoint",
+    )
+    parser.add_argument("--total_steps", type=int, default=200000)
+    parser.add_argument("--log_step", type=int, default=100)
+    parser.add_argument("--eval_step", type=int, default=1000)
+    parser.add_argument("--save_step", type=int, default=100)
+    parser.add_argument(
+        "--not_resume",
+        action="store_true",
+        help="Don't resume from the last checkpoint",
+    )
+
+    # for debugging
+    parser.add_argument("--limit_train_batches", type=int)
+    parser.add_argument("--limit_val_batches", type=int)
+    parser.add_argument("--fast_dev_run", action="store_true")
+
+    parser.add_argument("--backbone", type=str, default="XVector")
+    parser.add_argument("--pooling_type", type=str, default="TAP")
+    parser.add_argument("--loss_type", type=str, default="softmax")
+    parser.add_argument("--spk_embd_dim", type=int, default=1500)
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    logging.basicConfig()
+    logger.setLevel(logging.INFO)
+
+    args = parse_args()
+    voxceleb1 = Path(args.voxceleb1)
+    assert voxceleb1.is_dir()
+    save_to = Path(args.save_to)
+    save_to.mkdir(exist_ok=True, parents=True)
+
+    logger.info("Preparing preprocessor")
+    preprocessor = problem.Preprocessor(voxceleb1)
+
+    logger.info("Preparing train dataloader")
+    train_dataset = problem.TrainDataset(**preprocessor.train_data())
+    train_sampler = problem.TrainSampler(
+        train_dataset, max_timestamp=16000 * 1000, shuffle=True
+    )
+    train_sampler = DistributedBatchSamplerWrapper(
+        train_sampler, num_replicas=1, rank=0
+    )
+    train_dataloader = DataLoader(
+        train_dataset,
+        batch_sampler=train_sampler,
+        num_workers=6,
+        collate_fn=train_dataset.collate_fn,
+    )
+
+    logger.info("Preparing valid dataloader")
+    valid_dataset = problem.ValidDataset(
+        **preprocessor.valid_data(),
+        **train_dataset.statistics(),
+    )
+    valid_dataset.save_checkpoint(save_to / "valid_dataset.ckpt")
+    # valid_dataset_reload = Object.load_checkpoint(save_to / "valid_dataset.ckpt")
+    valid_sampler = problem.TrainSampler(
+        valid_dataset, max_timestamp=16000 * 1000, shuffle=True
+    )
+    valid_sampler = DistributedBatchSamplerWrapper(
+        valid_sampler, num_replicas=1, rank=0
+    )
+    valid_dataloader = DataLoader(
+        valid_dataset,
+        batch_sampler=valid_sampler,
+        num_workers=6,
+        collate_fn=valid_dataset.collate_fn,
+    )
+
+    logger.info("Preparing test dataloader")
+    test_dataset = problem.TestDataset(
+        **preprocessor.test_data(),
+        **train_dataset.statistics(),
+    )
+    test_dataset.save_checkpoint(save_to / "test_dataset.ckpt")
+    test_sampler = problem.TestSampler(test_dataset, 8)
+    test_sampler = DistributedBatchSamplerWrapper(test_sampler, num_replicas=1, rank=0)
+    test_dataloader = DataLoader(
+        test_dataset, batch_size=1, num_workers=6, collate_fn=test_dataset.collate_fn
+    )
+
+    upstream = S3PRLUpstream("apc")
+    # Have to specify the backbone, pooling_type, spk_embd_dim
+    downstream = problem.speaker_embedding_extractor(
+        backbone=args.backbone,
+        pooling_type=args.pooling_type,
+        input_size=upstream.output_size,
+        output_size=args.spk_embd_dim,
+    )
+    model = UpstreamDownstreamModel(upstream, downstream)
+    # Have to specify the loss_type
+    task = problem.Task(
+        model=model,
+        categories=preprocessor.statistics().category,
+        loss_type=args.loss_type,
+        trials=test_dataset.statistics().label,
+    )
+
+    optimizer = optim.Adam(task.parameters(), lr=1e-3)
+    lightning_task = LightningModuleSimpleWrapper(task, optimizer)
+
+    # The above is the usage of our library
+
+    # The below is pytorch-lightning specific usage, which can be very simple
+    # or very sophisticated, depending on how much you want to customized your
+    # training loop
+
+    checkpoint_callback = ModelCheckpoint(
+        dirpath=str(save_to),
+        filename="superb-sv-{step:02d}-{valid_0_accuracy:.2f}",
+        monitor="valid_0_accuracy",  # since might have multiple valid dataloaders
+        save_last=True,
+        save_top_k=3,  # top 3 best ckpt on valid
+        mode="max",  # higher, better
+        every_n_train_steps=args.save_step,
+    )
+
+    trainer = Trainer(
+        callbacks=[checkpoint_callback],
+        accelerator="gpu",
+        gpus=1,
+        max_steps=args.total_steps,
+        log_every_n_steps=args.log_step,
+        val_check_interval=args.eval_step,
+        limit_val_batches=args.limit_val_batches or 1.0,
+        limit_train_batches=args.limit_train_batches or 1.0,
+        fast_dev_run=args.fast_dev_run,
+    )
+
+    last_ckpt = save_to / "last.ckpt"
+    if args.not_resume or not last_ckpt.is_file():
+        last_ckpt = None
+
+    trainer.fit(
+        lightning_task,
+        train_dataloader,
+        val_dataloaders=valid_dataloader,
+        ckpt_path=last_ckpt,
+    )
+
+    trainer.test(
+        lightning_task,
+        dataloaders=test_dataloader,
+        ckpt_path=last_ckpt,
+    )
+
+
+if __name__ == "__main__":
+    main()

s3prl_s3prl_main/external_tools/install_espnet.sh

@@ -0,0 +1,20 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+if [ $# != 0 ]; then
+    echo "Usage: $0"
+    exit 1;
+fi
+
+if [ ! -e espnet.done ]; then
+    if ! python3 -c "import espnet2" &> /dev/null; then
+        pip install "espnet>=202308"
+    else
+        echo echo "espnet is already installed"
+    fi
+    touch espnet.done
+else
+    echo "espnet is already installed"
+fi
+
+

s3prl_s3prl_main/find_content.sh

@@ -0,0 +1,8 @@
+#!/bin/bash
+
+if [ -z $1 ]; then
+    echo "Usage: $0 <pattern>"
+    exit 1
+fi
+
+ag $1 $(cat valid_paths.txt)

s3prl_s3prl_main/hubconf.py

@@ -0,0 +1,4 @@
+from s3prl import hub
+
+for _option in hub.options():
+    globals()[_option] = getattr(hub, _option)

s3prl_s3prl_main/pyrightconfig.json

@@ -0,0 +1,5 @@
+{
+    "exclude": [
+        "**/result/**",
+    ],
+}

s3prl_s3prl_main/pytest.ini

@@ -0,0 +1,5 @@
+[pytest]
+minversion = 6.0
+addopts = -ra -q --log-cli-level INFO
+testpaths =
+    test

s3prl_s3prl_main/requirements/all.txt

@@ -0,0 +1,33 @@
+joblib >=0.12.4
+librosa >=0.7.2
+scipy >=1.5.4
+scikit-learn >=0.23.2
+pandas>=1.1.5
+matplotlib >=3.3.4
+Pillow >=6.2.2
+numba >=0.48
+cython >=0.29.21
+packaging >=20.9
+transformers >=4.10.0,<5.0
+dtw-python ==1.3.0
+asteroid ==0.4.4
+sacrebleu >=2.0.0
+h5py
+sox
+tabulate
+intervaltree
+lxml
+soundfile
+pysndfx
+nltk
+normalise
+editdistance
+easydict
+catalyst
+sentencepiece
+huggingface_hub >=0.2.1
+mutagen
+pydantic
+sed_eval
+more_itertools
+tensorboard

s3prl_s3prl_main/requirements/dev.txt

@@ -0,0 +1,11 @@
+python-dotenv
+pytest
+pytest-xdist
+tox
+pre-commit
+flake8 ==3.9.2
+black ==22.3.0
+isort ==5.10.1
+sphinx ==5.1.1  # follow https://github.com/sphinx-doc/sphinx/issues/6316
+furo
+importlib_metadata <5.0  # required by flake8