Merge pull request #65 from paganpasta/dev

Dev

README.md

@@ -48,7 +48,7 @@ Picking a model and doing a forward pass is as simple as ...
 
 ## What's New?
 
-- `FCN` and `DeepLabV3` added as new image segmentation models.
+- `FCN`, `DeepLabV3` and `LRASPP` added as new image segmentation models.
 - Backward incompatible changes to `v0.2.0` for loading a `pretrained` model.
 - Almost all image classification models are ported from `torchvision`.
 - New tutorial for generating `adversarial examples` and others coming soon.

docs/api/models/segmentation/lraspp.md

@@ -0,0 +1,12 @@
+# LRASPP
+
+
+::: eqxvision.models.LRASPP
+ selection:
+ members:
+ - __init__
+ - __call__
+
+---
+
+::: eqxvision.models.lraspp_mobilenet_v3_large

docs/index.md

@@ -36,8 +36,8 @@ pip install eqxvision
  ```
 
 ## What's New?
+- `FCN`, `DeepLabV3` and `LRASPP` segmentation models are now supported (checkout the [tutorial](getting_started/FCN_Segmentation.ipynb)).
 - Backward incompatible changes to `v0.2.0` for loading a `pretrained` model.
-- `FCN` and `DeepLabV3` segmentation models are now supported (checkout the [tutorial](getting_started/FCN_Segmentation.ipynb)).
 - Almost all image classification models are ported from `torchvision`.
 - New tutorial for generating [adversarial examples](getting_started/Adversarial_Attack.ipynb) and others coming soon.
 

eqxvision/__init__.py

@@ -1,4 +1,4 @@
 r"""Root package info."""
-__version__ = "0.2.5"
+__version__ = "0.2.6"
 
 from . import experimental, layers, models, utils

eqxvision/experimental.py

@@ -1,7 +1,7 @@
 from typing import Any, Callable
 
 import equinox as eqx
-from jaxtyping import PyTree
+import equinox.nn as nn
 
 
 class AuxData:
@@ -33,7 +33,7 @@ def __call__(self, x, *, key=None):
 
 
 def intermediate_layer_getter(
- model: PyTree, get_target_layers: Callable
+ model: "eqx.Module", get_target_layers: Callable
 ) -> "eqx.Module":
  """Wraps intermediate layers of a model for accessing intermediate activations. Based on a discussion
  [here](https://github.com/patrick-kidger/equinox/issues/186).
@@ -49,26 +49,34 @@ def intermediate_layer_getter(
  of layers from the `model`
 
  **Returns:**
- A `PyTree`, encapsulating `model` for storing intermediate outputs from target layers.
+ The returned model will now return a `tuple` with
 
- !!! info
- The returned model will now return a `tuple` with
+  0. The final output of `model`
+ 1. An ordered list of intermediate activations
 
- 1. The final output of `model`
- 2. An ordered list of intermediate activations
  """
  target_layers = get_target_layers(model)
  auxs, wrappers = zip(
  *[_make_intermediate_layer_wrapper() for _ in range(len(target_layers))]
  )
- model = eqx.tree_at(
- where=get_target_layers,
- pytree=model,
- replace=[
- wrapper(target_layer)
- for (wrapper, target_layer) in zip(wrappers, target_layers)
- ],
- )
+ if isinstance(model, nn.Sequential):
+ new_modules, updated_count = [], 0
+ for idx, module in enumerate(model.layers):
+ if idx in target_layers:
+ new_modules.append(wrappers[updated_count](module))
+ updated_count += 1
+ else:
+ new_modules.append(module)
+ model = nn.Sequential(new_modules)
+ else:
+ model = eqx.tree_at(
+ where=get_target_layers,
+ pytree=model,
+ replace=[
+ wrapper(target_layer)
+ for (wrapper, target_layer) in zip(wrappers, target_layers)
+ ],
+ )
 
  class IntermediateLayerGetter(eqx.Module):
  model: eqx.Module

eqxvision/models/__init__.py

@@ -102,3 +102,4 @@
 )
 from .segmentation.deeplabv3 import DeepLabV3, deeplabv3
 from .segmentation.fcn import FCN, fcn
+from .segmentation.lraspp import LRASPP, lraspp_mobilenet_v3_large

eqxvision/models/classification/mobilenetv3.py

@@ -361,7 +361,10 @@ def mobilenet_v3_large(torch_weights: str = None, **kwargs: Any) -> MobileNetV3:
 
  """
  arch = "mobilenet_v3_large"
- inverted_residual_setting, last_channel = _mobilenet_v3_conf(arch, **kwargs)
+ dilated = kwargs.pop("dilated", False)
+ inverted_residual_setting, last_channel = _mobilenet_v3_conf(
+ arch, dilated=dilated, **kwargs
+ )
  model = _mobilenet_v3(arch, inverted_residual_setting, last_channel, **kwargs)
  if torch_weights:
  model = load_torch_weights(model, torch_weights=torch_weights)

eqxvision/models/segmentation/__init__.py

@@ -1 +1 @@
-from . import deeplabv3, fcn
+from . import deeplabv3, fcn, lraspp

eqxvision/models/segmentation/deeplabv3.py

@@ -9,7 +9,7 @@
 from jaxtyping import Array
 
 from ...experimental import intermediate_layer_getter
-from ...utils import CLASSIFICATION_URLS, load_torch_weights
+from ...utils import load_torch_weights
 from ..classification import resnet
 from ._utils import _SimpleSegmentationModel
 from .fcn import FCNHead
@@ -139,17 +139,17 @@ def deeplabv3(
  num_classes: Optional[int] = 21,
  backbone: "eqx.Module" = None,
  intermediate_layers: Callable = None,
- classifier_module: "eqx.Module" = DeepLabHead,
+ classifier_module: "eqx.Module" = None,
  classifier_in_channels: int = 2048,
- aux_classifier_module: "eqx.Module" = FCNHead,
+ aux_classifier_module: "eqx.Module" = None,
  aux_in_channels: int = 1024,
  silence_layers: Callable = None,
  torch_weights: str = None,
  *,
  key: Optional["jax.random.PRNGKey"] = None,
 ) -> DeepLabV3:
  """Implements DeepLabV3 model from
- ["Rethinking Atrous Convolution for Semantic Image Segmentation"](https://arxiv.org/abs/1706.05587) paper.
+ [Rethinking Atrous Convolution for Semantic Image Segmentation](https://arxiv.org/abs/1706.05587) paper.
 
  !!! info "Sample call"
  ```python
@@ -167,7 +167,7 @@ def deeplabv3(
  - `num_classes`: Number of classes in the segmentation task.
  Also controls the final output shape `(num_classes, height, width)`. Defaults to `21`
  - `backbone`: The neural network to use for extracting features. If `None`, then all params are set to
- `DeepLabV3_RESNET50` with a **pre-trained** backbone but **untrained** DeepLabV3 heads
+ `DeepLabV3_RESNET50` with `untrained` weights
  - `intermediate_layers`: Layers from `backbone` to be used for generating output maps. Default sets it to
  `layer3` and `layer4` from `DeepLabV3_RESNET50`
  - `classifier_module`: Uses the `DeepLabHead` by default
@@ -179,15 +179,17 @@ def deeplabv3(
  the `fc` layers can be dropped. This is particularly useful when loading weights from `torchvision`. By
  default, `.fc` layer of a model is set to identity to avoid tracking weights.
  - `torch_weights`: A `Path` or `URL` for the `PyTorch` weights. Defaults to `None`
-
+ - `key`: A `jax.random.PRNGKey` used to provide randomness for parameter
  """
  if key is None:
  key = jr.PRNGKey(0)
  keys = jr.split(key, 2)
-
+ if not classifier_module:
+ classifier_module = DeepLabHead
+ if not aux_classifier_module:
+ aux_classifier_module = FCNHead
  if backbone is None:
  backbone = resnet.resnet50(
- torch_weights=CLASSIFICATION_URLS["resnet50"],
  replace_stride_with_dilation=[False, True, True],
  )
  num_layers = len(intermediate_layers(backbone))

eqxvision/models/segmentation/fcn.py

@@ -7,7 +7,7 @@
 import jax.random as jr
 
 from ...experimental import intermediate_layer_getter
-from ...utils import CLASSIFICATION_URLS, load_torch_weights
+from ...utils import load_torch_weights
 from ..classification import resnet
 from ._utils import _SimpleSegmentationModel
 
@@ -38,7 +38,7 @@ def fcn(
  num_classes: Optional[int] = 21,
  backbone: "eqx.Module" = None,
  intermediate_layers: Callable = None,
- classifier_module: "eqx.Module" = FCNHead,
+ classifier_module: "eqx.Module" = None,
  classifier_in_channels: int = 2048,
  aux_in_channels: int = None,
  silence_layers: Callable = None,
@@ -64,7 +64,7 @@ def fcn(
  - `num_classes`: Number of classes in the segmentation task.
  Also controls the final output shape `(num_classes, height, width)`. Defaults to `21`
  - `backbone`: The neural network to use for extracting features. If `None`, then all params are set to
- `FCN_RESNET50` with a **pre-trained** backbone but an **untrained** FCN
+ `FCN_RESNET50` with `untrained` weights
  - `intermediate_layers`: Layers from `backbone` to be used for generating output maps. Default sets it to
  `layer3` and `layer4` from `FCN_RESNET50`
  - `classifier_module`: Uses the `FCNHead` by default
@@ -75,15 +75,15 @@ def fcn(
  the `fc` layers can be dropped. This is particularly useful when loading weights from `torchvision`. By
  default, `.fc` layer of a model is set to identity to avoid tracking weights.
  - `torch_weights`: A `Path` or `URL` for the `PyTorch` weights. Defaults to `None`
-
+ - `key`: A `jax.random.PRNGKey` used to provide randomness for parameter
  """
  if key is None:
  key = jr.PRNGKey(0)
  keys = jr.split(key, 2)
-
+ if classifier_module is None:
+ classifier_module = FCNHead
  if backbone is None:
  backbone = resnet.resnet50(
- torch_weights=CLASSIFICATION_URLS["resnet50"],
  replace_stride_with_dilation=[False, True, True],
  )
  num_layers = len(intermediate_layers(backbone))