430 lines
14 KiB
Python
430 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.stochastic_depth import StochasticDepth
|
|
from torchvision.models._api import Weights, WeightsEnum
|
|
|
|
|
|
|
|
__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,
|
|
stochastic_depth_prob: float,
|
|
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=7, padding=3, 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 CNBlockConfig:
|
|
# Stores information listed at Section 3 of the ConvNeXt paper
|
|
def __init__(
|
|
self,
|
|
input_channels: int,
|
|
out_channels: Optional[int],
|
|
num_layers: int,
|
|
) -> None:
|
|
self.input_channels = input_channels
|
|
self.out_channels = out_channels
|
|
self.num_layers = num_layers
|
|
|
|
def __repr__(self) -> str:
|
|
s = self.__class__.__name__ + "("
|
|
s += "input_channels={input_channels}"
|
|
s += ", out_channels={out_channels}"
|
|
s += ", num_layers={num_layers}"
|
|
s += ")"
|
|
return s.format(**self.__dict__)
|
|
|
|
|
|
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,
|
|
input_features: int = 3,
|
|
device: str = 'cuda',
|
|
) -> 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.stages: List[nn.Module] = []
|
|
self.downsamples: List[nn.Module] = []
|
|
|
|
# Stem
|
|
firstconv_output_channels = block_setting[0].input_channels
|
|
self.stem = Conv2dNormActivation(
|
|
input_features,
|
|
firstconv_output_channels,
|
|
kernel_size=4,
|
|
stride=4,
|
|
padding=0,
|
|
norm_layer=norm_layer,
|
|
activation_layer=None,
|
|
bias=True,
|
|
).to(device)
|
|
|
|
layers_list: List[nn.Module] = []
|
|
|
|
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_list.append(nn.Sequential(*stage))
|
|
if cnf.out_channels is not None:
|
|
# Downsampling
|
|
layers_list.append(
|
|
nn.Sequential(
|
|
norm_layer(cnf.input_channels),
|
|
nn.Conv2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
|
|
).to(device)
|
|
)
|
|
|
|
self.layers = nn.Sequential(*layers_list).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.stem(x)
|
|
features: List[Tensor] = []
|
|
x = self.layers(x)
|
|
x = self.act(x)
|
|
|
|
return x, features
|
|
|
|
def forward(self, x: Tensor) -> Tensor:
|
|
return self._forward_impl(x)
|
|
|
|
|
|
|
|
class ConvNeXtReverseDepthNormals(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=False,
|
|
) -> 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)
|
|
|
|
# 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
|
|
layers_list: 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))
|
|
stage_block_id += 1
|
|
layers_list.append(nn.Sequential(*stage).to(device))
|
|
else:
|
|
layers_list.append(nn.Identity().to(device))
|
|
|
|
if cnf.out_channels is not None:
|
|
# Upsampling
|
|
layers_list.append(
|
|
nn.Sequential(
|
|
norm_layer(cnf.input_channels),
|
|
nn.ConvTranspose2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
|
|
).to(device)
|
|
)
|
|
|
|
self.layers = nn.Sequential(*layers_list).to(device)
|
|
|
|
for cnf in block_setting[-1:]:
|
|
# Bottlenecks
|
|
stage: List[nn.Module] = []
|
|
for _ in range(cnf.num_layers):
|
|
stage.append(block(cnf.input_channels, layer_scale, 0))
|
|
self.stage_normals = nn.Sequential(*stage).to(device)
|
|
|
|
self.last = nn.Conv2d(
|
|
block_setting[-1].input_channels,
|
|
1,
|
|
kernel_size=3,
|
|
padding='same',
|
|
).to(device)
|
|
self.last_normals = nn.Conv2d(
|
|
block_setting[-1].input_channels,
|
|
3,
|
|
kernel_size=3,
|
|
padding='same',
|
|
).to(device)
|
|
# 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:
|
|
if self.use_skip:
|
|
input, features = x
|
|
else:
|
|
input = x
|
|
x = self.first(input)
|
|
for i in range(5 - 2):
|
|
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)
|
|
f = self.upsamples[-1](x)
|
|
|
|
x = self.stages[-1](f)
|
|
depth = self.last(x)
|
|
|
|
x = self.stage_normals(f)
|
|
normals = self.last_normals(x)
|
|
|
|
return depth, normals
|
|
|
|
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,
|
|
device: str = 'cuda',
|
|
use_skip=False,
|
|
) -> 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.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
|
|
layers_list: 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))
|
|
stage_block_id += 1
|
|
layers_list.append(nn.Sequential(*stage).to(device))
|
|
else:
|
|
layers_list.append(nn.Identity().to(device))
|
|
|
|
if cnf.out_channels is not None:
|
|
# Upnsampling
|
|
layers_list.append(
|
|
nn.Sequential(
|
|
norm_layer(cnf.input_channels),
|
|
nn.ConvTranspose2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
|
|
).to(device)
|
|
)
|
|
|
|
self.layers = nn.Sequential(*layers_list).to(device)
|
|
self.last = nn.Conv2d(
|
|
block_setting[-1].input_channels,
|
|
1,
|
|
kernel_size=3,
|
|
padding='same',
|
|
).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:
|
|
if self.use_skip:
|
|
raise NotImplementedError
|
|
input, features = x
|
|
else:
|
|
input = x
|
|
x = self.first(input)
|
|
x = self.layers(x)
|
|
out = self.last(x)
|
|
|
|
return out
|
|
|
|
def forward(self, x) -> Tensor:
|
|
return self._forward_impl(x)
|
|
|
|
|
|
def _convnext(
|
|
block_setting: List[CNBlockConfig],
|
|
stochastic_depth_prob: float,
|
|
weights: Optional[WeightsEnum],
|
|
reversed: bool,
|
|
input_features: int,
|
|
**kwargs: Any,
|
|
) -> ConvNeXtBase:
|
|
|
|
if reversed:
|
|
model = ConvNeXtReverseBase(block_setting, stochastic_depth_prob=stochastic_depth_prob,**kwargs)
|
|
else:
|
|
model = ConvNeXtBase(block_setting, stochastic_depth_prob=stochastic_depth_prob, input_features=input_features, **kwargs)
|
|
|
|
if weights is not None:
|
|
model.load_state_dict(weights.get_state_dict())
|
|
|
|
return model
|
|
|
|
|
|
class ConvNeXt(nn.Module):
|
|
def __init__(self, block_setting, stochastic_depth_prob=0.0, reversed=False, input_features=3, **kwargs) -> None:
|
|
super().__init__()
|
|
|
|
self.model = _convnext(block_setting, stochastic_depth_prob, None, reversed, input_features, **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 |