File size: 3,504 Bytes
62762da | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | # Copyright (c) Meta Platforms, Inc. and affiliates.
# Copyright 2025 AI for Oncology Research Group. All Rights Reserved.
#
# 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.
#
# References:
# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
from typing import Callable, Optional
import torch
from torch import nn
class Mlp(nn.Module):
"""Multi-layer perceptron (MLP) module.
Creates a simple MLP with two linear layers and an activation function in between and dropout after each layer.
Parameters
----------
in_features : int
Number of input features.
hidden_features : int, optional
Number of hidden features, by default 4 * in_features.
out_features : int, optional
Number of output features, by default in_features.
act_layer : Callable[..., nn.Module], optional
Activation layer, by default nn.GELU.
drop : float, optional
Dropout rate, by default 0.0.
bias : bool, optional
Whether to use bias in the linear layers, by default True.
"""
def __init__(
self,
in_features: int,
hidden_features: Optional[int] = None,
out_features: Optional[int] = None,
act_layer: Callable[..., nn.Module] = nn.GELU,
drop: float = 0.0,
bias: bool = True,
) -> None:
"""Inits :class:`Mlp`.
Parameters
----------
in_features : int
Number of input features.
hidden_features : int, optional
Number of hidden features, by default 4 * in_features.
out_features : int, optional
Number of output features, by default in_features.
act_layer : Callable[..., nn.Module], optional
Activation layer, by default nn.GELU.
drop : float, optional
Dropout rate, by default 0.0.
bias : bool, optional
Whether to use bias in the linear layers, by default True.
"""
super().__init__()
out_features = out_features or in_features
hidden_features = hidden_features or in_features
self.fc1 = nn.Linear(in_features, hidden_features, bias=bias)
self.act = act_layer()
self.fc2 = nn.Linear(hidden_features, out_features, bias=bias)
self.drop = nn.Dropout(drop)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""Forward pass of :class:`Mlp`.
Parameters
----------
x : torch.Tensor
Input tensor of shape (B, N, C) where B is the batch size, N is the sequence length, and C is
the feature dimension.
Returns
-------
torch.Tensor
Output tensor of shape (B, N, out_features) after applying the MLP.
"""
x = self.fc1(x)
x = self.act(x)
x = self.drop(x)
x = self.fc2(x)
x = self.drop(x)
return x
|