|
|
import torch |
|
|
from typing import Callable, List, Optional |
|
|
|
|
|
|
|
|
class MLP(torch.nn.Sequential): |
|
|
"""This block implements the multi-layer perceptron (MLP) module. |
|
|
Adapted for backward compatibility from the torchvision library: |
|
|
https://pytorch.org/vision/0.14/generated/torchvision.ops.MLP.html |
|
|
|
|
|
LICENSE: |
|
|
|
|
|
From PyTorch: |
|
|
|
|
|
Copyright (c) 2016- Facebook, Inc (Adam Paszke) |
|
|
Copyright (c) 2014- Facebook, Inc (Soumith Chintala) |
|
|
Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert) |
|
|
Copyright (c) 2012-2014 Deepmind Technologies (Koray Kavukcuoglu) |
|
|
Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu) |
|
|
Copyright (c) 2011-2013 NYU (Clement Farabet) |
|
|
Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston) |
|
|
Copyright (c) 2006 Idiap Research Institute (Samy Bengio) |
|
|
Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz) |
|
|
|
|
|
From Caffe2: |
|
|
|
|
|
Copyright (c) 2016-present, Facebook Inc. All rights reserved. |
|
|
|
|
|
All contributions by Facebook: |
|
|
Copyright (c) 2016 Facebook Inc. |
|
|
|
|
|
All contributions by Google: |
|
|
Copyright (c) 2015 Google Inc. |
|
|
All rights reserved. |
|
|
|
|
|
All contributions by Yangqing Jia: |
|
|
Copyright (c) 2015 Yangqing Jia |
|
|
All rights reserved. |
|
|
|
|
|
All contributions by Kakao Brain: |
|
|
Copyright 2019-2020 Kakao Brain |
|
|
|
|
|
All contributions by Cruise LLC: |
|
|
Copyright (c) 2022 Cruise LLC. |
|
|
All rights reserved. |
|
|
|
|
|
All contributions from Caffe: |
|
|
Copyright(c) 2013, 2014, 2015, the respective contributors |
|
|
All rights reserved. |
|
|
|
|
|
All other contributions: |
|
|
Copyright(c) 2015, 2016 the respective contributors |
|
|
All rights reserved. |
|
|
|
|
|
Caffe2 uses a copyright model similar to Caffe: each contributor holds |
|
|
copyright over their contributions to Caffe2. The project versioning records |
|
|
all such contribution and copyright details. If a contributor wants to further |
|
|
mark their specific copyright on a particular contribution, they should |
|
|
indicate their copyright solely in the commit message of the change when it is |
|
|
committed. |
|
|
|
|
|
All rights reserved. |
|
|
|
|
|
Redistribution and use in source and binary forms, with or without |
|
|
modification, are permitted provided that the following conditions are met: |
|
|
|
|
|
1. Redistributions of source code must retain the above copyright |
|
|
notice, this list of conditions and the following disclaimer. |
|
|
|
|
|
2. Redistributions in binary form must reproduce the above copyright |
|
|
notice, this list of conditions and the following disclaimer in the |
|
|
documentation and/or other materials provided with the distribution. |
|
|
|
|
|
3. Neither the names of Facebook, Deepmind Technologies, NYU, NEC Laboratories America |
|
|
and IDIAP Research Institute nor the names of its contributors may be |
|
|
used to endorse or promote products derived from this software without |
|
|
specific prior written permission. |
|
|
|
|
|
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" |
|
|
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE |
|
|
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE |
|
|
ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE |
|
|
LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR |
|
|
CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF |
|
|
SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS |
|
|
INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN |
|
|
CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) |
|
|
ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE |
|
|
POSSIBILITY OF SUCH DAMAGE. |
|
|
|
|
|
|
|
|
Args: |
|
|
in_channels (int): Number of channels of the input |
|
|
hidden_channels (List[int]): List of the hidden channel dimensions |
|
|
norm_layer (Callable[..., torch.nn.Module], optional): Norm layer that will be stacked on top of the linear layer. If ``None`` this layer won't be used. Default: ``None`` |
|
|
activation_layer (Callable[..., torch.nn.Module], optional): Activation function which will be stacked on top of the normalization layer (if not None), otherwise on top of the linear layer. If ``None`` this layer won't be used. Default: ``torch.nn.ReLU`` |
|
|
inplace (bool, optional): Parameter for the activation layer, which can optionally do the operation in-place. |
|
|
Default is ``None``, which uses the respective default values of the ``activation_layer`` and Dropout layer. |
|
|
bias (bool): Whether to use bias in the linear layer. Default ``True`` |
|
|
dropout (float): The probability for the dropout layer. Default: 0.0 |
|
|
""" |
|
|
|
|
|
def __init__( |
|
|
self, |
|
|
in_channels: int, |
|
|
hidden_channels: List[int], |
|
|
activation_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.ReLU, |
|
|
inplace: Optional[bool] = None, |
|
|
bias: bool = True, |
|
|
dropout: float = 0.0, |
|
|
): |
|
|
params = {} if inplace is None else {"inplace": inplace} |
|
|
|
|
|
layers = [] |
|
|
in_dim = in_channels |
|
|
for hidden_dim in hidden_channels[:-1]: |
|
|
layers.append(torch.nn.Linear(in_dim, hidden_dim, bias=bias)) |
|
|
layers.append(activation_layer(**params)) |
|
|
layers.append(torch.nn.Dropout(dropout, **params)) |
|
|
in_dim = hidden_dim |
|
|
|
|
|
layers.append(torch.nn.Linear(in_dim, hidden_channels[-1], bias=bias)) |
|
|
layers.append(torch.nn.Dropout(dropout, **params)) |
|
|
|
|
|
super().__init__(*layers) |
|
|
|
|
|
|
|
|
def batch_idx(x: torch.Tensor, idx: torch.Tensor): |
|
|
assert ( |
|
|
idx.shape == x.shape[: len(idx.shape)] |
|
|
), f"Index shape must match a prefix of x.shape. Got idx: {idx.shape} and x: {x.shape}" |
|
|
|
|
|
remaining_shape = x.shape[len(idx.shape) + 1 :] |
|
|
x = x.flatten(start_dim=idx.dim() + 1) |
|
|
indices = [] |
|
|
for i in range(len(idx.shape)): |
|
|
raw = torch.arange(idx.shape[i], device=idx.device) |
|
|
desired_shape = [1] * len(idx.shape) |
|
|
desired_shape[i] = idx.shape[i] |
|
|
indices.append(raw.reshape(desired_shape)) |
|
|
indices.append(idx) |
|
|
result = x[tuple(indices)] |
|
|
return result.reshape(idx.shape + remaining_shape) |
|
|
|
