tap-ct-b-2d / patch_embed.py

TimVeenboer

model commit

62762da 10 months ago

11.6 kB

# 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

from .helpers import make_2tuple, make_3tuple


class PatchEmbed(nn.Module):
    """Patch embedding layer for Vision Transformers for 2D images.
    This layer divides the input image into patches and projects them into a higher-dimensional space.
    Parameters
    ----------
    img_size : int or tuple[int, int], optional
        Size of the input image. If an integer is provided, it is assumed to be square (img_size, img_size).
        If a tuple is provided, it should be of the form (height, width), by default 224.
    patch_size : int or tuple[int, int], optional
        Size of the patches to be extracted from the input image. If an integer is provided, it is assumed to be square
        (patch_size, patch_size). If a tuple is provided, it should be of the form (height, width), by default 16.
    in_chans : int, optional
        Number of input channels in the image, by default 3 (for RGB images).
    embed_dim : int, optional
        Dimension of the embedding space to which the patches will be projected, by default 768.
    norm_layer : Callable, optional
        Normalization layer to apply to the embeddings, by default None. If None, no normalization is applied.
    flatten_embedding : bool, optional
        Whether to flatten the embedding output, by default True.
    """

    def __init__(
        self,
        img_size: int | tuple[int, int] = 224,
        patch_size: int | tuple[int, int] = 16,
        in_chans: int = 3,
        embed_dim: int = 768,
        norm_layer: Optional[Callable] = None,
        flatten_embedding: bool = True,
    ) -> None:
        """Inits :class:`PatchEmbed`.
        Parameters
        ----------
        img_size : int or tuple[int, int], optional
            Size of the input image. If an integer is provided, it is assumed to be square (img_size, img_size).
            If a tuple is provided, it should be of the form (height, width), by default 224.
        patch_size : int or tuple[int, int], optional
            Size of the patches to be extracted from the input image. If an integer is provided, it is assumed to be square
            (patch_size, patch_size). If a tuple is provided, it should be of the form (height, width), by default 16.
        in_chans : int, optional
            Number of input channels in the image, by default 3 (for RGB images).
        embed_dim : int, optional
            Dimension of the embedding space to which the patches will be projected, by default 768.
        norm_layer : Callable, optional
            Normalization layer to apply to the embeddings, by default None. If None, no normalization is applied.
        flatten_embedding : bool, optional
            Whether to flatten the embedding output, by default True.
        """
        super().__init__()

        image_HW = make_2tuple(img_size)
        patch_HW = make_2tuple(patch_size)
        patch_grid_size = (
            image_HW[0] // patch_HW[0],
            image_HW[1] // patch_HW[1],
        )

        self.img_size = image_HW
        self.patch_size = patch_HW
        self.patches_resolution = patch_grid_size
        self.num_patches = patch_grid_size[0] * patch_grid_size[1]

        self.in_chans = in_chans
        self.embed_dim = embed_dim

        self.flatten_embedding = flatten_embedding

        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_HW, stride=patch_HW)
        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass of :class:`PatchEmbed`.
        Parameters
        ----------
        x : torch.Tensor
            Input tensor of shape (B, C, H, W) where B is the batch size, C is the number of channels,
            H is the height, and W is the width of the input image.
        Raises
        ------
        ValueError
            If the input image dimensions are not compatible with the patch size.
        """
        _, _, H, W = x.shape
        patch_H, patch_W = self.patch_size
        if H % patch_H != 0:
            raise ValueError(f"Input image height {H} is not a multiple of patch height {patch_H}")
        if W % patch_W != 0:
            raise ValueError(f"Input image width {W} is not a multiple of patch width: {patch_W}")

        x = self.proj(x)  # B C H W
        H, W = x.size(2), x.size(3)
        x = x.flatten(2).transpose(1, 2)  # B HW C

        x = self.norm(x)
        if not self.flatten_embedding:
            x = x.reshape(-1, H, W, self.embed_dim)  # B H W C
        return x

    def flops(self) -> float:
        """Calculate the number of floating point operations (FLOPs) for the patch embedding layer.
        Returns
        -------
        float
            The number of FLOPs for the patch embedding layer.
        """
        Ho, Wo = self.patches_resolution
        flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1])
        if not isinstance(self.norm, nn.Identity):
            flops += Ho * Wo * self.embed_dim
        return flops


class PatchEmbed3d(nn.Module):
    """Patch embedding layer for Vision Transformers for 3D images.
    This layer divides the input 3D image volume into patches and projects them into a higher-dimensional space.
    Parameters
    ----------
    img_size : int or tuple[int, int, int], optional
        Size of the input image volume. If an integer is provided, it is assumed to be cubic (img_size, img_size, img_size).
        If a tuple is provided, it should be of the form (depth, height, width), by default 224.
    patch_size : int or tuple[int, int, int], optional
        Size of the patches to be extracted from the input image volume. If an integer is provided, it is assumed to be cubic
        (patch_size, patch_size, patch_size). If a tuple is provided, it should be of the form (depth, height, width), by default 16.
    in_chans : int, optional
        Number of input channels in the image volume, by default 3 (for RGB images).
    embed_dim : int, optional
        Dimension of the embedding space to which the patches will be projected, by default 768.
    norm_layer : Callable, optional
        Normalization layer to apply to the embeddings, by default None. If None, no normalization is applied.
    flatten_embedding : bool, optional
        Whether to flatten the embedding output, by default True.
    """

    def __init__(
        self,
        img_size: int | tuple[int, int, int] = 224,
        patch_size: int | tuple[int, int, int] = 16,
        in_chans: int = 3,
        embed_dim: int = 768,
        norm_layer: Optional[Callable] = None,
        flatten_embedding: bool = True,
    ) -> None:
        """Inits :class:`PatchEmbed3d`.
        Parameters
        ----------
        img_size : int or tuple[int, int, int], optional
            Size of the input image volume. If an integer is provided, it is assumed to be cubic
            (img_size, img_size, img_size).
            If a tuple is provided, it should be of the form (depth, height, width), by default 224.
        patch_size : int or tuple[int, int, int], optional
            Size of the patches to be extracted from the input image volume. If an integer is provided, it is
            assumed to be cubic (patch_size, patch_size, patch_size). If a tuple is provided, it should be of the
            form (depth, height, width), by default 16.
        in_chans : int, optional
            Number of input channels in the image volume, by default 3 (for RGB images).
        embed_dim : int, optional
            Dimension of the embedding space to which the patches will be projected, by default 768.
        norm_layer : Callable, optional
            Normalization layer to apply to the embeddings, by default None. If None, no normalization is applied.
        flatten_embedding : bool, optional
            Whether to flatten the embedding output, by default True.
        """
        super().__init__()

        image_DHW = make_3tuple(img_size)
        patch_DHW = make_3tuple(patch_size)

        patch_grid_size = (
            image_DHW[0] // patch_DHW[0],
            image_DHW[1] // patch_DHW[1],
            image_DHW[2] // patch_DHW[2],
        )

        self.img_size = image_DHW
        self.patch_size = patch_DHW
        self.patches_resolution = patch_grid_size
        self.num_patches = patch_grid_size[0] * patch_grid_size[1] * patch_grid_size[2]

        self.in_chans = in_chans
        self.embed_dim = embed_dim

        self.flatten_embedding = flatten_embedding
        self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_DHW, stride=patch_DHW)
        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass of :class:`PatchEmbed3d`.
        Parameters
        ----------
        x : torch.Tensor
            Input tensor of shape (B, C, D, H, W) where B is the batch size, C is the number of channels,
            D is the depth, H is the height, and W is the width of the input volume.
        Raises
        ------
        ValueError
            If the input volume dimensions are not compatible with the patch size.
        """
        _, _, D, H, W = x.shape
        patch_D, patch_H, patch_W = self.patch_size
        if D % patch_D != 0:
            raise ValueError(f"Input volume depth {D} is not a multiple of patch depth {patch_D}")
        if H % patch_H != 0:
            raise ValueError(f"Input volume height {H} is not a multiple of patch height {patch_H}")
        if W % patch_W != 0:
            raise ValueError(f"Input volume width {W} is not a multiple of patch width {patch_W}")

        x = self.proj(x)  # B C D H W
        D, H, W = x.size(2), x.size(3), x.size(4)
        x = x.flatten(2).transpose(1, 2)  # B (DHW) C

        x = self.norm(x)
        if not self.flatten_embedding:
            x = x.reshape(-1, D, H, W, self.embed_dim)  # B D H W C
        return x

    def flops(self) -> float:
        """Calculate the number of floating point operations (FLOPs) for the patch embedding 3D layer.
        Returns
        -------
        float
            The number of FLOPs for the patch embedding layer.
        """
        Do, Ho, Wo = self.patches_resolution
        flops = (
            Do
            * Ho
            * Wo
            * self.embed_dim
            * self.in_chans
            * (self.patch_size[0] * self.patch_size[1] * self.patch_size[2])
        )
        if not isinstance(self.norm, nn.Identity):
            flops += Do * Ho * Wo * self.embed_dim
        return flops

	# 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

	from .helpers import make_2tuple, make_3tuple


	class PatchEmbed(nn.Module):
	"""Patch embedding layer for Vision Transformers for 2D images.

	This layer divides the input image into patches and projects them into a higher-dimensional space.

	Parameters
	----------
	img_size : int or tuple[int, int], optional
	Size of the input image. If an integer is provided, it is assumed to be square (img_size, img_size).
	If a tuple is provided, it should be of the form (height, width), by default 224.
	patch_size : int or tuple[int, int], optional
	Size of the patches to be extracted from the input image. If an integer is provided, it is assumed to be square
	(patch_size, patch_size). If a tuple is provided, it should be of the form (height, width), by default 16.
	in_chans : int, optional
	Number of input channels in the image, by default 3 (for RGB images).
	embed_dim : int, optional
	Dimension of the embedding space to which the patches will be projected, by default 768.
	norm_layer : Callable, optional
	Normalization layer to apply to the embeddings, by default None. If None, no normalization is applied.
	flatten_embedding : bool, optional
	Whether to flatten the embedding output, by default True.
	"""

	def __init__(
	self,
	img_size: int \| tuple[int, int] = 224,
	patch_size: int \| tuple[int, int] = 16,
	in_chans: int = 3,
	embed_dim: int = 768,
	norm_layer: Optional[Callable] = None,
	flatten_embedding: bool = True,
	) -> None:
	"""Inits :class:`PatchEmbed`.

	Parameters
	----------
	img_size : int or tuple[int, int], optional
	Size of the input image. If an integer is provided, it is assumed to be square (img_size, img_size).
	If a tuple is provided, it should be of the form (height, width), by default 224.
	patch_size : int or tuple[int, int], optional
	Size of the patches to be extracted from the input image. If an integer is provided, it is assumed to be square
	(patch_size, patch_size). If a tuple is provided, it should be of the form (height, width), by default 16.
	in_chans : int, optional
	Number of input channels in the image, by default 3 (for RGB images).
	embed_dim : int, optional
	Dimension of the embedding space to which the patches will be projected, by default 768.
	norm_layer : Callable, optional
	Normalization layer to apply to the embeddings, by default None. If None, no normalization is applied.
	flatten_embedding : bool, optional
	Whether to flatten the embedding output, by default True.
	"""
	super().__init__()

	image_HW = make_2tuple(img_size)
	patch_HW = make_2tuple(patch_size)
	patch_grid_size = (
	image_HW[0] // patch_HW[0],
	image_HW[1] // patch_HW[1],
	)

	self.img_size = image_HW
	self.patch_size = patch_HW
	self.patches_resolution = patch_grid_size
	self.num_patches = patch_grid_size[0] * patch_grid_size[1]

	self.in_chans = in_chans
	self.embed_dim = embed_dim

	self.flatten_embedding = flatten_embedding

	self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_HW, stride=patch_HW)
	self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	"""Forward pass of :class:`PatchEmbed`.

	Parameters
	----------
	x : torch.Tensor
	Input tensor of shape (B, C, H, W) where B is the batch size, C is the number of channels,
	H is the height, and W is the width of the input image.

	Raises
	------
	ValueError
	If the input image dimensions are not compatible with the patch size.
	"""
	_, _, H, W = x.shape
	patch_H, patch_W = self.patch_size
	if H % patch_H != 0:
	raise ValueError(f"Input image height {H} is not a multiple of patch height {patch_H}")
	if W % patch_W != 0:
	raise ValueError(f"Input image width {W} is not a multiple of patch width: {patch_W}")

	x = self.proj(x) # B C H W
	H, W = x.size(2), x.size(3)
	x = x.flatten(2).transpose(1, 2) # B HW C

	x = self.norm(x)
	if not self.flatten_embedding:
	x = x.reshape(-1, H, W, self.embed_dim) # B H W C
	return x

	def flops(self) -> float:
	"""Calculate the number of floating point operations (FLOPs) for the patch embedding layer.

	Returns
	-------
	float
	The number of FLOPs for the patch embedding layer.
	"""
	Ho, Wo = self.patches_resolution
	flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1])
	if not isinstance(self.norm, nn.Identity):
	flops += Ho * Wo * self.embed_dim
	return flops


	class PatchEmbed3d(nn.Module):
	"""Patch embedding layer for Vision Transformers for 3D images.

	This layer divides the input 3D image volume into patches and projects them into a higher-dimensional space.

	Parameters
	----------
	img_size : int or tuple[int, int, int], optional
	Size of the input image volume. If an integer is provided, it is assumed to be cubic (img_size, img_size, img_size).
	If a tuple is provided, it should be of the form (depth, height, width), by default 224.
	patch_size : int or tuple[int, int, int], optional
	Size of the patches to be extracted from the input image volume. If an integer is provided, it is assumed to be cubic
	(patch_size, patch_size, patch_size). If a tuple is provided, it should be of the form (depth, height, width), by default 16.
	in_chans : int, optional
	Number of input channels in the image volume, by default 3 (for RGB images).
	embed_dim : int, optional
	Dimension of the embedding space to which the patches will be projected, by default 768.
	norm_layer : Callable, optional
	Normalization layer to apply to the embeddings, by default None. If None, no normalization is applied.
	flatten_embedding : bool, optional
	Whether to flatten the embedding output, by default True.
	"""

	def __init__(
	self,
	img_size: int \| tuple[int, int, int] = 224,
	patch_size: int \| tuple[int, int, int] = 16,
	in_chans: int = 3,
	embed_dim: int = 768,
	norm_layer: Optional[Callable] = None,
	flatten_embedding: bool = True,
	) -> None:
	"""Inits :class:`PatchEmbed3d`.

	Parameters
	----------
	img_size : int or tuple[int, int, int], optional
	Size of the input image volume. If an integer is provided, it is assumed to be cubic
	(img_size, img_size, img_size).
	If a tuple is provided, it should be of the form (depth, height, width), by default 224.
	patch_size : int or tuple[int, int, int], optional
	Size of the patches to be extracted from the input image volume. If an integer is provided, it is
	assumed to be cubic (patch_size, patch_size, patch_size). If a tuple is provided, it should be of the
	form (depth, height, width), by default 16.
	in_chans : int, optional
	Number of input channels in the image volume, by default 3 (for RGB images).
	embed_dim : int, optional
	Dimension of the embedding space to which the patches will be projected, by default 768.
	norm_layer : Callable, optional
	Normalization layer to apply to the embeddings, by default None. If None, no normalization is applied.
	flatten_embedding : bool, optional
	Whether to flatten the embedding output, by default True.
	"""
	super().__init__()

	image_DHW = make_3tuple(img_size)
	patch_DHW = make_3tuple(patch_size)

	patch_grid_size = (
	image_DHW[0] // patch_DHW[0],
	image_DHW[1] // patch_DHW[1],
	image_DHW[2] // patch_DHW[2],
	)

	self.img_size = image_DHW
	self.patch_size = patch_DHW
	self.patches_resolution = patch_grid_size
	self.num_patches = patch_grid_size[0] * patch_grid_size[1] * patch_grid_size[2]

	self.in_chans = in_chans
	self.embed_dim = embed_dim

	self.flatten_embedding = flatten_embedding
	self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_DHW, stride=patch_DHW)
	self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	"""Forward pass of :class:`PatchEmbed3d`.

	Parameters
	----------
	x : torch.Tensor
	Input tensor of shape (B, C, D, H, W) where B is the batch size, C is the number of channels,
	D is the depth, H is the height, and W is the width of the input volume.

	Raises
	------
	ValueError
	If the input volume dimensions are not compatible with the patch size.
	"""
	_, _, D, H, W = x.shape
	patch_D, patch_H, patch_W = self.patch_size
	if D % patch_D != 0:
	raise ValueError(f"Input volume depth {D} is not a multiple of patch depth {patch_D}")
	if H % patch_H != 0:
	raise ValueError(f"Input volume height {H} is not a multiple of patch height {patch_H}")
	if W % patch_W != 0:
	raise ValueError(f"Input volume width {W} is not a multiple of patch width {patch_W}")

	x = self.proj(x) # B C D H W
	D, H, W = x.size(2), x.size(3), x.size(4)
	x = x.flatten(2).transpose(1, 2) # B (DHW) C

	x = self.norm(x)
	if not self.flatten_embedding:
	x = x.reshape(-1, D, H, W, self.embed_dim) # B D H W C
	return x

	def flops(self) -> float:
	"""Calculate the number of floating point operations (FLOPs) for the patch embedding 3D layer.

	Returns
	-------
	float
	The number of FLOPs for the patch embedding layer.
	"""
	Do, Ho, Wo = self.patches_resolution
	flops = (
	Do
	* Ho
	* Wo
	* self.embed_dim
	* self.in_chans
	* (self.patch_size[0] * self.patch_size[1] * self.patch_size[2])
	)
	if not isinstance(self.norm, nn.Identity):
	flops += Do * Ho * Wo * self.embed_dim
	return flops