Fix broken paths from standalone model reorganization

Author: albertfgu

CHANGELOG.md

@@ -54,14 +54,9 @@ To address differences between models trained on earlier versions and the curren
 #### Models and Features
 - Updated version of S4 module, including new measures and theory from [[How to Train Your HiPPO](https://arxiv.org/abs/2206.12037)] (https://github.com/HazyResearch/state-spaces/issues/21, https://github.com/HazyResearch/state-spaces/issues/54)
 - Complete version of S4D module from [[On the Parameterization and Initialization of Diagonal State Space Models](https://arxiv.org/abs/2206.11893)]
-- [State forwarding](src/models/s4/README.md#state-forwarding) (https://github.com/HazyResearch/state-spaces/issues/49, https://github.com/HazyResearch/state-spaces/issues/56)
-- Support for S4 variants including DSS and GSS ([documentation](src/models/s4/README.md#other-variants))
+- [State forwarding](models/s4/README.md#state-forwarding) (https://github.com/HazyResearch/state-spaces/issues/49, https://github.com/HazyResearch/state-spaces/issues/56)
+- Support for S4 variants including DSS and GSS ([documentation](models/s4/README.md#other-variants))
 
-<!--
-####  Compilation of additional resources
-  - Recommended resources for understanding S4-style models, including the [Simplifying S4 blog](https://hazyresearch.stanford.edu/blog/2022-06-11-simplifying-s4) ([code](https://github.com/HazyResearch/state-spaces/tree/simple/src/models/sequence/ss/s4_simple)) and a minimal pedagogical version of S4D ([code](src/models/s4/s4d.py))
-  - Tips & Tricks page for getting started with tuning S4
--->
 
 #### Bug fixes and library compatibility issues
 - PyTorch 1.11 had a [Dropout bug](https://github.com/pytorch/pytorch/issues/77081) which is now avoided with a custom Dropout implementation (https://github.com/HazyResearch/state-spaces/issues/42, https://github.com/HazyResearch/state-spaces/issues/22)

README.md

@@ -68,7 +68,7 @@ Installation usually works out of the box with `pip install pykeops cmake` which
 
 ### S4 Module
 
-Self-contained files for the S4 layer and variants can be found in [src/models/s4/](./src/models/s4/),
+Self-contained files for the S4 layer and variants can be found in [models/s4/](./models/s4/),
 which includes instructions for calling the module.
 
 See [notebooks/](notebooks/) for visualizations explaining some concepts behind HiPPO and S4.
@@ -97,7 +97,7 @@ One important feature of this codebase is supporting parameters that require dif
 In particular, the SSM kernel is particularly sensitive to the $(A, B)$ (and sometimes $\Delta$ parameters),
 so the learning rate on these parameters is sometimes lowered and the weight decay is always set to $0$.
 
-See the method `register` in the model (e.g. [s4d.py](src/models/s4/s4d.py)) and the function `setup_optimizer` in the training script (e.g. [example.py](example.py)) for an examples of how to implement this in external repos.
+See the method `register` in the model (e.g. [s4d.py](py)) and the function `setup_optimizer` in the training script (e.g. [example.py](example.py)) for an examples of how to implement this in external repos.
 
 <!--
 Our logic for setting these parameters can be found in the `OptimModule` class under `src/models/sequence/ss/kernel.py` and the corresponding optimizer hook in `SequenceLightningModule.configure_optimizers` under `train.py`

example.py

@@ -3,7 +3,7 @@
 This code borrows heavily from https://github.com/kuangliu/pytorch-cifar.
 
 This file only depends on the standalone S4 layer
-available in src/models/s4/
+available in /models/s4/
 
 * Train standard sequential CIFAR:
     python -m example

models/sashimi/sashimi.py

@@ -2,7 +2,7 @@
 SaShiMi backbone.
 
 Use this backbone in your own models. You'll also need to copy over the
-standalone S4 layer, which can be found at `state-spaces/src/models/s4/`
+standalone S4 layer, which can be found at `state-spaces/models/s4/`
 
 It's Raw! Audio Generation with State-Space Models
 Karan Goel, Albert Gu, Chris Donahue, Christopher Re.
@@ -16,7 +16,7 @@
 
 from einops import rearrange
 
-from src.models.s4.s4 import LinearActivation, S4
+from models.s4.s4 import LinearActivation, S4Block as S4
 
 class DownPool(nn.Module):
     def __init__(self, d_input, expand, pool):

src/utils/registry.py

@@ -36,7 +36,7 @@
     "model":                 "src.models.sequence.backbones.model.SequenceModel",
     "unet":                  "src.models.sequence.backbones.unet.SequenceUNet",
     "sashimi":               "src.models.sequence.backbones.sashimi.Sashimi",
-    "sashimi_standalone":    "sashimi.sashimi.Sashimi",
+    "sashimi_standalone":    "models.sashimi.sashimi.Sashimi",
     # Baseline RNNs
     "lstm":                  "src.models.baselines.lstm.TorchLSTM",
     "gru":                   "src.models.baselines.gru.TorchGRU",
@@ -46,12 +46,8 @@
     "stackedrnn":            "src.models.baselines.samplernn.StackedRNN",
     "stackedrnn_baseline":   "src.models.baselines.samplernn.StackedRNNBaseline",
     "samplernn":             "src.models.baselines.samplernn.SampleRNN",
-    "dcgru":                 "src.models.baselines.dcgru.DCRNNModel_classification",
-    "dcgru_ss":              "src.models.baselines.dcgru.DCRNNModel_nextTimePred",
     # Baseline CNNs
     "ckconv":                "src.models.baselines.ckconv.ClassificationCKCNN",
-    "wavegan":               "src.models.baselines.wavegan.WaveGANDiscriminator", # DEPRECATED
-    "denseinception":        "src.models.baselines.dense_inception.DenseInception",
     "wavenet":               "src.models.baselines.wavenet.WaveNetModel",
     "torch/resnet2d":        "src.models.baselines.resnet.TorchVisionResnet",  # 2D ResNet
     # Nonaka 1D CNN baselines
@@ -69,15 +65,13 @@
     "timm/convnext_micro":   "src.models.baselines.convnext_timm.convnext_micro",
     "timm/resnet50":         "src.models.baselines.resnet_timm.resnet50", # Can also register many other variants in resnet_timm
     "timm/convnext_tiny_3d": "src.models.baselines.convnext_timm.convnext3d_tiny",
-    # Segmentation models
-    "convnext_unet_tiny":    "src.models.segmentation.convnext_unet.convnext_tiny_unet",
 }
 
 layer = {
     "id":         "src.models.sequence.base.SequenceIdentity",
     "lstm":       "src.models.baselines.lstm.TorchLSTM",
     "standalone": "models.s4.s4.S4Block",
-    "s4d":        "src.models.s4.s4d.S4D",
+    "s4d":        "models.s4.s4d.S4D",
     "ffn":        "src.models.sequence.modules.ffn.FFN",
     "sru":        "src.models.sequence.rnns.sru.SRURNN",
     "rnn":        "src.models.sequence.rnns.rnn.RNN",  # General RNN wrapper
@@ -90,8 +84,6 @@
     "s4":         "src.models.sequence.modules.s4block.S4Block",
     "s4nd":       "src.models.sequence.modules.s4nd.S4ND",
     "mega":       "src.models.sequence.modules.mega.MegaBlock",
-    "h3":         "src.models.sequence.experimental.h3.H3",
-    "h4":         "src.models.sequence.experimental.h4.H4",
     # 'packedrnn': 'models.sequence.rnns.packedrnn.PackedRNN',
 }