pose encoder model
This commit is contained in:
+74
-84
@@ -94,6 +94,7 @@ class ConvNeXtBase(nn.Module):
|
||||
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__()
|
||||
@@ -115,7 +116,7 @@ class ConvNeXtBase(nn.Module):
|
||||
# Stem
|
||||
firstconv_output_channels = block_setting[0].input_channels
|
||||
self.stem = Conv2dNormActivation(
|
||||
3,
|
||||
input_features,
|
||||
firstconv_output_channels,
|
||||
kernel_size=4,
|
||||
stride=4,
|
||||
@@ -125,6 +126,7 @@ class ConvNeXtBase(nn.Module):
|
||||
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
|
||||
@@ -136,16 +138,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))
|
||||
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.layers = nn.Sequential(*layers_list).to(device)
|
||||
self.act = nn.Tanh()
|
||||
|
||||
for m in self.modules():
|
||||
@@ -157,11 +160,7 @@ class ConvNeXtBase(nn.Module):
|
||||
def _forward_impl(self, x: Tensor) -> Tensor:
|
||||
x = self.stem(x)
|
||||
features: List[Tensor] = []
|
||||
for i in range(4 - 1):
|
||||
x = self.stages[i](x)
|
||||
features.append(x)
|
||||
x = self.downsamples[i](x)
|
||||
x = self.stages[-1](x)
|
||||
x = self.layers(x)
|
||||
x = self.act(x)
|
||||
|
||||
return x, features
|
||||
@@ -171,7 +170,7 @@ class ConvNeXtBase(nn.Module):
|
||||
|
||||
|
||||
|
||||
class ConvNeXtReverseBase2(nn.Module):
|
||||
class ConvNeXtReverseDepthNormals(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
block_setting: List[CNBlockConfig],
|
||||
@@ -217,8 +216,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
|
||||
@@ -228,19 +227,21 @@ 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))
|
||||
stage_block_id += 1
|
||||
self.stages.append(nn.Sequential(*stage).to(device))
|
||||
layers_list.append(nn.Sequential(*stage).to(device))
|
||||
else:
|
||||
self.stages.append(nn.Identity().to(device))
|
||||
layers_list.append(nn.Identity().to(device))
|
||||
|
||||
if cnf.out_channels is not None:
|
||||
# Upnsampling
|
||||
self.upsamples.append(
|
||||
# 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] = []
|
||||
@@ -275,7 +276,10 @@ class ConvNeXtReverseBase2(nn.Module):
|
||||
|
||||
|
||||
def _forward_impl(self, x) -> Tensor:
|
||||
input, features = x
|
||||
if self.use_skip:
|
||||
input, features = x
|
||||
else:
|
||||
input = x
|
||||
x = self.first(input)
|
||||
for i in range(5 - 2):
|
||||
if self.use_skip:
|
||||
@@ -306,10 +310,14 @@ class ConvNeXtReverseBase(nn.Module):
|
||||
stochastic_depth_prob: float = 0.0,
|
||||
layer_scale: float = 1e-6,
|
||||
block: Optional[Callable[..., nn.Module]] = None,
|
||||
norm_layer: 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])):
|
||||
@@ -321,71 +329,44 @@ class ConvNeXtReverseBase(nn.Module):
|
||||
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,
|
||||
),
|
||||
)
|
||||
)
|
||||
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
|
||||
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))
|
||||
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.append(
|
||||
# Upnsampling
|
||||
layers_list.append(
|
||||
nn.Sequential(
|
||||
norm_layer(cnf.out_channels),
|
||||
nn.ConvTranspose2d(cnf.out_channels, cnf.input_channels, kernel_size=2, stride=2),
|
||||
)
|
||||
norm_layer(cnf.input_channels),
|
||||
nn.ConvTranspose2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
|
||||
).to(device)
|
||||
)
|
||||
|
||||
self.features = nn.Sequential(*reversed(layers))
|
||||
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)):
|
||||
@@ -393,12 +374,20 @@ class ConvNeXtReverseBase(nn.Module):
|
||||
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:
|
||||
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)
|
||||
|
||||
|
||||
@@ -407,13 +396,14 @@ def _convnext(
|
||||
stochastic_depth_prob: float,
|
||||
weights: Optional[WeightsEnum],
|
||||
reversed: bool,
|
||||
input_features: int,
|
||||
**kwargs: Any,
|
||||
) -> ConvNeXtBase:
|
||||
|
||||
if reversed:
|
||||
model = ConvNeXtReverseBase2(block_setting, stochastic_depth_prob=stochastic_depth_prob, **kwargs)
|
||||
model = ConvNeXtReverseBase(block_setting, stochastic_depth_prob=stochastic_depth_prob,**kwargs)
|
||||
else:
|
||||
model = ConvNeXtBase(block_setting, stochastic_depth_prob=stochastic_depth_prob, **kwargs)
|
||||
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())
|
||||
@@ -422,10 +412,10 @@ def _convnext(
|
||||
|
||||
|
||||
class ConvNeXt(nn.Module):
|
||||
def __init__(self, block_setting, stochastic_depth_prob=0.0, reversed=False, **kwargs) -> None:
|
||||
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, **kwargs)
|
||||
self.model = _convnext(block_setting, stochastic_depth_prob, None, reversed, input_features, **kwargs)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
x = self.model(x)
|
||||
|
||||
Reference in New Issue
Block a user