161 lines
6.0 KiB
Python
161 lines
6.0 KiB
Python
from typing import List, Dict, Any, Tuple
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
class ConvSingleLayer(nn.Module):
|
|
"""A single convolutional layer followed by an optional ReLU activation and optional downsampling.
|
|
|
|
Shape is tensor of shape `(batch_size, in_channels, height, width)`.
|
|
|
|
Args:
|
|
in_channels (int): Number of input channels.
|
|
out_channels (int): Number of output channels.
|
|
kernel_size (int): Size of the convolutional kernel.
|
|
padding (int): Padding size for convolution.
|
|
downsample (bool): If True, applies a 2x2 MaxPooling layer.
|
|
apply_activation (bool, optional): If True, applies ReLU activation. Defaults to True.
|
|
|
|
pragma: skip_doc
|
|
"""
|
|
|
|
@classmethod # TODO: Dummy method to suppress ComfyUI warning. DO NOT REMOVE
|
|
def INPUT_TYPES(cls) -> Dict[str, Any]:
|
|
return {}
|
|
|
|
def __init__(self,
|
|
in_channels: int,
|
|
out_channels: int,
|
|
kernel_size: int,
|
|
padding: int,
|
|
downsample: bool,
|
|
bias:bool = True,
|
|
apply_activation: bool = True) -> None:
|
|
super().__init__()
|
|
self.conv = nn.Conv2d(in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
kernel_size=kernel_size,
|
|
bias=bias,
|
|
padding=padding)
|
|
self.apply_activation = apply_activation
|
|
self.activation = nn.ReLU()
|
|
self.downsample = downsample
|
|
self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)
|
|
|
|
if self.apply_activation:
|
|
torch.nn.init.kaiming_normal_(self.conv.weight,
|
|
nonlinearity='relu')
|
|
else:
|
|
torch.nn.init.xavier_normal_(self.conv.weight)
|
|
|
|
|
|
def forward(self, inputs: torch.Tensor) -> torch.Tensor:
|
|
"""Forward pass of the convolutional layer.
|
|
|
|
Args:
|
|
inputs (torch.Tensor): Input tensor of shape `(batch_size, in_channels, height, width)`.
|
|
|
|
Returns:
|
|
torch.Tensor: Output tensor after convolution, optional activation, and optional pooling.
|
|
"""
|
|
x = self.conv(inputs)
|
|
if self.downsample:
|
|
x = self.maxpool(x)
|
|
if self.apply_activation:
|
|
x = self.activation(x)
|
|
return x
|
|
|
|
|
|
class ConvModel(nn.Module):
|
|
"""A multi-layer convolutional neural network with ReLU activation after each convolutional layer, except for the final dense layer.
|
|
|
|
Args:
|
|
input_dim (Tuple[int, int, int]): Shape of a single input sample (C, H, W).
|
|
penultimate_dim (int): Number of features before the final dense layer. If 0 is specified, this value is automatically computed internally.
|
|
output_dim (int): Number of output features.
|
|
channel_list (List[int]): List of channel sizes for each layer. This should not include the inputs channel.
|
|
Input channel is taken from the first element of input_dim.
|
|
kernel_size_list (List[int]): List of kernel sizes for each convolutional layer.
|
|
padding_list (List[int]): List of padding values for each convolutional layer.
|
|
downsample_list (List[bool]): List indicating whether downsampling should be applied after each layer.
|
|
|
|
pragma: skip_doc
|
|
"""
|
|
|
|
@classmethod # TODO: Dummy method to suppress ComfyUI warning
|
|
def INPUT_TYPES(cls) -> Dict[str, Any]:
|
|
return {}
|
|
|
|
def __init__(self,
|
|
input_dim: Tuple[int, int, int],
|
|
penultimate_dim: int,
|
|
output_dim: int,
|
|
channel_list: List[int],
|
|
kernel_size_list: List[int],
|
|
padding_list: List[int],
|
|
downsample_list: List[bool]) -> None:
|
|
super().__init__()
|
|
|
|
# Validate list lengths
|
|
if (len(channel_list) != len(kernel_size_list) or
|
|
len(kernel_size_list) != len(padding_list) or
|
|
len(padding_list) != len(downsample_list)):
|
|
raise ValueError("channel_list, kernel_size_list, padding_list, and downsample_list must have the same length.")
|
|
|
|
self.channel_list = [input_dim[0]] + channel_list
|
|
|
|
# Initialize convolutional layers
|
|
layers = []
|
|
for i in range(len(self.channel_list) - 1):
|
|
layers.append(
|
|
ConvSingleLayer(
|
|
in_channels=self.channel_list[i],
|
|
out_channels=self.channel_list[i + 1],
|
|
kernel_size=kernel_size_list[i],
|
|
padding=padding_list[i],
|
|
downsample=downsample_list[i]
|
|
)
|
|
)
|
|
|
|
self.conv_layers = nn.ModuleList(layers)
|
|
|
|
if penultimate_dim == 0:
|
|
# Try to auto compute penultimate_dim by processing the dummy input
|
|
with torch.no_grad():
|
|
dummy_input = torch.randn(1, *input_dim) # B, C, H, W
|
|
x = dummy_input
|
|
for layer in self.conv_layers:
|
|
x = layer(x)
|
|
checked_penultimate_dim = x.numel() # Number of elements
|
|
else:
|
|
checked_penultimate_dim = penultimate_dim
|
|
|
|
self.final_dense = nn.Linear(checked_penultimate_dim, output_dim, bias=True)
|
|
|
|
def forward(self, inputs: torch.Tensor) -> torch.Tensor:
|
|
"""Forward pass of the model.
|
|
|
|
Args:
|
|
inputs (torch.Tensor): Input tensor of shape `(batch_size, C, H, W)`.
|
|
|
|
Returns:
|
|
torch.Tensor: Output tensor of shape `(batch_size, output_dim)`.
|
|
|
|
Raises:
|
|
ValueError: If inputs is not a rank 4 tensor.
|
|
"""
|
|
if inputs.dim() not in [3, 4]:
|
|
raise ValueError("Inputs is not a rank 3 or 4 tensor.")
|
|
|
|
if inputs.dim() == 3:
|
|
x = torch.unsqueeze(inputs, 1)
|
|
else:
|
|
x = inputs
|
|
|
|
for layer in self.conv_layers:
|
|
x = layer(x)
|
|
|
|
x = torch.flatten(x, start_dim=1) # Flatten except batch-axis
|
|
x = self.final_dense(x)
|
|
return x
|