luni/cvpr_2023
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix convenxt layers

Browse Source
This commit is contained in:
luni
2023-04-06 17:24:52 +00:00
parent c0d180e72f
commit 51c8b1e069
1 changed files with 18 additions and 183 deletions
Download Patch File Download Diff File Expand all files Collapse all files
models/convnext.py
+18 -183
View File
@@ -8,25 +8,10 @@ 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)
@@ -91,11 +76,7 @@ class ConvNeXtBase(nn.Module):
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] = []
layers_list: List[nn.Module] = []
# Stem
firstconv_output_channels = block_setting[0].input_channels
@@ -121,17 +102,17 @@ class ConvNeXtBase(nn.Module):
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))
layers_list.append(nn.Sequential(*stage))
if cnf.out_channels is not None:
# Downsampling
self.downsamples.append(
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.act = nn.Tanh()
self.layers = nn.Sequential(*layers_list).to(device)
for m in self.modules():
if isinstance(m, (nn.Conv2d, nn.Linear)):
@@ -140,25 +121,15 @@ class ConvNeXtBase(nn.Module):
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)
x = self.layers(x)
return x
def forward(self, x: Tensor) -> Tensor:
return self._forward_impl(x)
class ConvNeXtReverseBase2(nn.Module):
class ConvNeXtReverseBase(nn.Module):
def __init__(
self,
block_setting: List[CNBlockConfig],
@@ -187,21 +158,6 @@ class ConvNeXtReverseBase2(nn.Module):
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),
@@ -210,8 +166,8 @@ class ConvNeXtReverseBase2(nn.Module):
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] = []
layers_list: List[nn.Module] = []
for cnf in block_setting[1:]:
if cnf.input_channels is not None:
# Bottlenecks
@@ -221,33 +177,30 @@ class ConvNeXtReverseBase2(nn.Module):
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))
layers_list.append(nn.Sequential(*stage))
else:
self.stages.append(nn.Identity().to(device))
layers_list.append(nn.Identity())
if cnf.out_channels is not None:
# Upsampling
self.upsamples.append(
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,
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)):
@@ -257,134 +210,16 @@ class ConvNeXtReverseBase2(nn.Module):
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.layers(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,
@@ -397,7 +232,7 @@ def _convnext(
) -> 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)
model = ConvNeXtReverseBase(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)