424 lines
14 KiB
Python
424 lines
14 KiB
Python
# taken from torchvision.models.convnext
|
|
|
|
from functools import partial
|
|
from typing import Any, Callable, List, Optional, Sequence
|
|
|
|
import torch
|
|
from torch import nn, Tensor
|
|
from torch.nn import functional as F
|
|
|
|
from torchvision.ops.misc import Conv2dNormActivation, Permute
|
|
# from torchvision.ops.misc import ConvNormActivation as Conv2dNormActivation
|
|
# from torchvision.ops import Permute
|
|
from torchvision.ops.stochastic_depth import StochasticDepth
|
|
from utils.convnext_utils import CNBlockConfig
|
|
|
|
|
|
__all__ = [
|
|
"ConvNeXt",
|
|
"ConvNeXt_Tiny_Weights",
|
|
"ConvNeXt_Small_Weights",
|
|
"ConvNeXt_Base_Weights",
|
|
"ConvNeXt_Large_Weights",
|
|
"convnext_tiny",
|
|
"convnext_small",
|
|
"convnext_base",
|
|
"convnext_large",
|
|
]
|
|
|
|
|
|
class LayerNorm2d(nn.LayerNorm):
|
|
def forward(self, x: Tensor) -> Tensor:
|
|
x = x.permute(0, 2, 3, 1)
|
|
x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
|
|
x = x.permute(0, 3, 1, 2)
|
|
return x
|
|
|
|
|
|
class CNBlock(nn.Module):
|
|
def __init__(
|
|
self,
|
|
dim,
|
|
layer_scale: float = 1e-6,
|
|
stochastic_depth_prob: float = 0.0,
|
|
kernel_size: int = 7,
|
|
norm_layer: Optional[Callable[..., nn.Module]] = None,
|
|
) -> None:
|
|
super().__init__()
|
|
if norm_layer is None:
|
|
norm_layer = partial(nn.LayerNorm, eps=1e-6)
|
|
|
|
self.block = nn.Sequential(
|
|
nn.Conv2d(dim, dim, kernel_size=kernel_size, padding='same', groups=dim, bias=True),
|
|
Permute([0, 2, 3, 1]),
|
|
norm_layer(dim),
|
|
nn.Linear(in_features=dim, out_features=4 * dim, bias=True),
|
|
nn.GELU(),
|
|
nn.Linear(in_features=4 * dim, out_features=dim, bias=True),
|
|
Permute([0, 3, 1, 2]),
|
|
)
|
|
self.layer_scale = nn.Parameter(torch.ones(dim, 1, 1) * layer_scale)
|
|
self.stochastic_depth = StochasticDepth(stochastic_depth_prob, "row")
|
|
|
|
def forward(self, input: Tensor) -> Tensor:
|
|
result = self.layer_scale * self.block(input)
|
|
result = self.stochastic_depth(result)
|
|
result += input
|
|
return result
|
|
|
|
|
|
class ConvNeXtBase(nn.Module):
|
|
def __init__(
|
|
self,
|
|
block_setting: List[CNBlockConfig],
|
|
stochastic_depth_prob: float = 0.0,
|
|
layer_scale: float = 1e-6,
|
|
block: Optional[Callable[..., nn.Module]] = None,
|
|
norm_layer: Optional[Callable[..., nn.Module]] = None,
|
|
device: str = 'cuda',
|
|
kernel_size: int = 7,
|
|
) -> None:
|
|
super().__init__()
|
|
|
|
if not block_setting:
|
|
raise ValueError("The block_setting should not be empty")
|
|
elif not (isinstance(block_setting, Sequence) and all([isinstance(s, CNBlockConfig) for s in block_setting])):
|
|
raise TypeError("The block_setting should be List[CNBlockConfig]")
|
|
|
|
if block is None:
|
|
block = CNBlock
|
|
|
|
if norm_layer is None:
|
|
norm_layer = partial(LayerNorm2d, eps=1e-6)
|
|
|
|
# self.permute_cf = Permute([0, 3, 1, 2])
|
|
# self.permute_cl = Permute([0, 2, 3, 1])
|
|
|
|
self.stages: List[nn.Module] = []
|
|
self.downsamples: List[nn.Module] = []
|
|
|
|
# Stem
|
|
firstconv_output_channels = block_setting[0].input_channels
|
|
self.stem = Conv2dNormActivation(
|
|
3,
|
|
firstconv_output_channels,
|
|
kernel_size=4,
|
|
stride=4,
|
|
padding=0,
|
|
norm_layer=norm_layer,
|
|
activation_layer=None,
|
|
bias=True,
|
|
).to(device)
|
|
|
|
|
|
total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
|
|
stage_block_id = 0
|
|
for cnf in block_setting:
|
|
# Bottlenecks
|
|
stage: List[nn.Module] = []
|
|
for _ in range(cnf.num_layers):
|
|
# adjust stochastic depth probability based on the depth of the stage block
|
|
sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
|
|
stage.append(block(cnf.input_channels, layer_scale, sd_prob, kernel_size))
|
|
stage_block_id += 1
|
|
self.stages.append(nn.Sequential(*stage).to(device))
|
|
if cnf.out_channels is not None:
|
|
# Downsampling
|
|
self.downsamples.append(
|
|
nn.Sequential(
|
|
norm_layer(cnf.input_channels),
|
|
nn.Conv2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
|
|
).to(device)
|
|
)
|
|
|
|
self.act = nn.Tanh()
|
|
|
|
for m in self.modules():
|
|
if isinstance(m, (nn.Conv2d, nn.Linear)):
|
|
nn.init.trunc_normal_(m.weight, std=0.02)
|
|
if m.bias is not None:
|
|
nn.init.zeros_(m.bias)
|
|
|
|
def _forward_impl(self, x: Tensor) -> Tensor:
|
|
# x = self.permute_cf(x)
|
|
x = self.stem(x)
|
|
# features: List[Tensor] = []
|
|
for i in range(len(self.stages) - 1):
|
|
x = self.stages[i](x)
|
|
# features.append(x)
|
|
x = self.downsamples[i](x)
|
|
x = self.stages[-1](x)
|
|
x = self.act(x)
|
|
|
|
# x = self.permute_cl(x)
|
|
return x
|
|
|
|
def forward(self, x: Tensor) -> Tensor:
|
|
return self._forward_impl(x)
|
|
|
|
|
|
|
|
class ConvNeXtReverseBase2(nn.Module):
|
|
def __init__(
|
|
self,
|
|
block_setting: List[CNBlockConfig],
|
|
stochastic_depth_prob: float = 0.0,
|
|
layer_scale: float = 1e-6,
|
|
block: Optional[Callable[..., nn.Module]] = None,
|
|
norm_layer: Optional[Callable[..., nn.Module]] = None,
|
|
device: str = 'cuda',
|
|
use_skip: bool = False,
|
|
is_label: bool = False,
|
|
nr_of_labels: Optional[int] = None,
|
|
kernel_size: int = 7,
|
|
) -> None:
|
|
super().__init__()
|
|
|
|
self.use_skip = use_skip
|
|
|
|
if not block_setting:
|
|
raise ValueError("The block_setting should not be empty")
|
|
elif not (isinstance(block_setting, Sequence) and all([isinstance(s, CNBlockConfig) for s in block_setting])):
|
|
raise TypeError("The block_setting should be List[CNBlockConfig]")
|
|
|
|
if block is None:
|
|
block = CNBlock
|
|
|
|
if norm_layer is None:
|
|
norm_layer = partial(LayerNorm2d, eps=1e-6)
|
|
|
|
# self.permute_cf = Permute([0, 3, 1, 2])
|
|
# self.permute_cl = Permute([0, 2, 3, 1])
|
|
|
|
# Stem
|
|
# firstconv_output_channels = block_setting[0].input_channels
|
|
# self.first = nn.Sequential(
|
|
# nn.ConvTranspose2d(
|
|
# firstconv_output_channels,
|
|
# firstconv_output_channels,
|
|
# kernel_size=2,
|
|
# stride=2,
|
|
# padding=0,
|
|
# ),
|
|
# norm_layer(firstconv_output_channels),
|
|
# ).to(device)
|
|
self.first = nn.Sequential(
|
|
nn.ConvTranspose2d(block_setting[0].input_channels, block_setting[0].out_channels, kernel_size=2, stride=2),
|
|
norm_layer(block_setting[0].out_channels),
|
|
).to(device)
|
|
|
|
|
|
total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
|
|
stage_block_id = 0
|
|
self.stages: List[nn.Module] = []
|
|
self.upsamples: List[nn.Module] = []
|
|
for cnf in block_setting[1:]:
|
|
if cnf.input_channels is not None:
|
|
# Bottlenecks
|
|
stage: List[nn.Module] = []
|
|
for _ in range(cnf.num_layers):
|
|
# adjust stochastic depth probability based on the depth of the stage block
|
|
sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
|
|
stage.append(block(cnf.input_channels, layer_scale, sd_prob, kernel_size))
|
|
stage_block_id += 1
|
|
self.stages.append(nn.Sequential(*stage).to(device))
|
|
else:
|
|
self.stages.append(nn.Identity().to(device))
|
|
|
|
if cnf.out_channels is not None:
|
|
# Upsampling
|
|
self.upsamples.append(
|
|
nn.Sequential(
|
|
norm_layer(cnf.input_channels),
|
|
nn.ConvTranspose2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
|
|
).to(device)
|
|
)
|
|
|
|
self.last = nn.Conv2d(
|
|
block_setting[-1].input_channels,
|
|
3 if not is_label else nr_of_labels,
|
|
kernel_size=3,
|
|
padding='same',
|
|
).to(device)
|
|
self.activation = nn.Tanh() if not is_label else nn.Softmax(dim=1)
|
|
|
|
# self.last = nn.Sequential(
|
|
# Permute([0, 2, 3, 1]),
|
|
# nn.Linear(in_features=block_setting[-1].input_channels, out_features=1, bias=True),
|
|
# Permute([0, 3, 1, 2]),
|
|
# ).to(device)
|
|
|
|
|
|
for m in self.modules():
|
|
if isinstance(m, (nn.Conv2d, nn.Linear)):
|
|
nn.init.trunc_normal_(m.weight, std=0.02)
|
|
if m.bias is not None:
|
|
nn.init.zeros_(m.bias)
|
|
|
|
|
|
def _forward_impl(self, x) -> Tensor:
|
|
# input, features = x
|
|
# x = self.permute_cf(x)
|
|
|
|
x = self.first(x)
|
|
for i in range(len(self.stages)):
|
|
# if self.use_skip:
|
|
# x = self.stages[i](torch.concat((x, features[2-i]), dim=1))
|
|
# else:
|
|
x = self.stages[i](x)
|
|
x = self.upsamples[i](x)
|
|
|
|
# x = self.stages[-2](x)
|
|
# x = self.upsamples[-1](x)
|
|
|
|
# x = self.stages[-1](x)
|
|
x = self.last(x)
|
|
x = self.activation(x)
|
|
# x = self.permute_cl(x)
|
|
|
|
return x
|
|
|
|
def forward(self, x) -> Tensor:
|
|
return self._forward_impl(x)
|
|
|
|
|
|
class ConvNeXtReverseBase(nn.Module):
|
|
def __init__(
|
|
self,
|
|
block_setting: List[CNBlockConfig],
|
|
stochastic_depth_prob: float = 0.0,
|
|
layer_scale: float = 1e-6,
|
|
block: Optional[Callable[..., nn.Module]] = None,
|
|
norm_layer: Optional[Callable[..., nn.Module]] = None
|
|
) -> None:
|
|
super().__init__()
|
|
|
|
if not block_setting:
|
|
raise ValueError("The block_setting should not be empty")
|
|
elif not (isinstance(block_setting, Sequence) and all([isinstance(s, CNBlockConfig) for s in block_setting])):
|
|
raise TypeError("The block_setting should be List[CNBlockConfig]")
|
|
|
|
if block is None:
|
|
block = CNBlock
|
|
|
|
if norm_layer is None:
|
|
norm_layer = partial(LayerNorm2d, eps=1e-6)
|
|
|
|
layers: List[nn.Module] = []
|
|
|
|
|
|
# add bottleneck which will be executed at the end, to process after the final larger upsampling
|
|
|
|
|
|
|
|
# Stem
|
|
firstconv_output_channels = block_setting[0].input_channels
|
|
|
|
layers.append(
|
|
nn.Sequential(
|
|
nn.Conv2d(16, 1, kernel_size=7, padding=3, bias=True),
|
|
|
|
)
|
|
)
|
|
|
|
stage: List[nn.Module] = []
|
|
for _ in range(3):
|
|
stage.append(block(16, layer_scale, 0.0))
|
|
layers.append(nn.Sequential(*stage))
|
|
|
|
layers.append(
|
|
# Convtr NormActivation(
|
|
# firstconv_output_channels,
|
|
# 3,
|
|
# kernel_size=4,
|
|
# stride=4,
|
|
# padding=0,
|
|
# norm_layer=None,
|
|
# activation_layer=None,
|
|
# bias=True,
|
|
# )
|
|
nn.Sequential(
|
|
nn.ConvTranspose2d(
|
|
firstconv_output_channels,
|
|
16,
|
|
kernel_size=4,
|
|
stride=4,
|
|
padding=0,
|
|
),
|
|
)
|
|
)
|
|
|
|
total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
|
|
stage_block_id = 0
|
|
for cnf in block_setting:
|
|
# Bottlenecks
|
|
stage: List[nn.Module] = []
|
|
for _ in range(cnf.num_layers):
|
|
# adjust stochastic depth probability based on the depth of the stage block
|
|
sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
|
|
stage.append(block(cnf.input_channels, layer_scale, sd_prob))
|
|
stage_block_id += 1
|
|
layers.append(nn.Sequential(*stage))
|
|
if cnf.out_channels is not None:
|
|
# Upsampling
|
|
layers.append(
|
|
nn.Sequential(
|
|
norm_layer(cnf.out_channels),
|
|
nn.ConvTranspose2d(cnf.out_channels, cnf.input_channels, kernel_size=2, stride=2),
|
|
)
|
|
)
|
|
|
|
self.features = nn.Sequential(*reversed(layers))
|
|
|
|
for m in self.modules():
|
|
if isinstance(m, (nn.Conv2d, nn.Linear)):
|
|
nn.init.trunc_normal_(m.weight, std=0.02)
|
|
if m.bias is not None:
|
|
nn.init.zeros_(m.bias)
|
|
|
|
def _forward_impl(self, x: Tensor) -> Tensor:
|
|
x = self.features(x)
|
|
# x = torch.tanh(x)
|
|
return x
|
|
|
|
def forward(self, x: Tensor) -> Tensor:
|
|
return self._forward_impl(x)
|
|
|
|
|
|
def _convnext(
|
|
block_setting: List[CNBlockConfig],
|
|
stochastic_depth_prob: float,
|
|
reversed: bool,
|
|
is_label: bool,
|
|
nr_of_labels: Optional[int],
|
|
kernel_size: int,
|
|
device='cuda',
|
|
**kwargs: Any,
|
|
) -> ConvNeXtBase:
|
|
|
|
if reversed:
|
|
model = ConvNeXtReverseBase2(block_setting, stochastic_depth_prob=stochastic_depth_prob, is_label=is_label, nr_of_labels=nr_of_labels, kernel_size=kernel_size, device=device, **kwargs)
|
|
else:
|
|
model = ConvNeXtBase(block_setting, stochastic_depth_prob=stochastic_depth_prob, kernel_size=kernel_size, device=device, **kwargs)
|
|
|
|
|
|
return model
|
|
|
|
|
|
class ConvNeXt(nn.Module):
|
|
def __init__(self, block_setting, stochastic_depth_prob=0.0, reversed=False, is_label=False, nr_of_labels=None, kernel_size=7, device='cuda', **kwargs) -> None:
|
|
super().__init__()
|
|
|
|
self.model = _convnext(block_setting, stochastic_depth_prob, reversed, is_label, nr_of_labels, kernel_size, device=device, **kwargs)
|
|
|
|
def forward(self, x: Tensor) -> Tensor:
|
|
x = self.model(x)
|
|
return x
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
model = ConvNeXt().to('cuda')
|
|
test = model(torch.zeros((4, 3, 512, 256)).to('cuda'))
|
|
|
|
pass |