pose autoencoder pipeline
This commit is contained in:
@@ -0,0 +1,183 @@
|
||||
import os
|
||||
import sys
|
||||
|
||||
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
|
||||
|
||||
import torch
|
||||
from torch import nn, Tensor
|
||||
from torch.nn import functional as F
|
||||
from functools import partial
|
||||
from typing import Any, Callable, List, Optional, Sequence
|
||||
|
||||
torch.set_float32_matmul_precision('high')
|
||||
|
||||
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
|
||||
from models.convnext import ConvNeXt, CNBlockConfig, CNBlock, LayerNorm2d
|
||||
|
||||
|
||||
class ConvNeXtReversePose(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)
|
||||
|
||||
stage: List[nn.Module] = []
|
||||
for _ in range(3):
|
||||
stage.append(block(block_setting[0].input_channels, layer_scale, 0))
|
||||
self.stage_smooth = nn.Sequential(*stage).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[:-1]).to(device)
|
||||
self.stage_pose = nn.Sequential(*layers_list[-1:]).to(device)
|
||||
|
||||
stage: List[nn.Module] = []
|
||||
for _ in range(cnf.num_layers):
|
||||
stage.append(block(block_setting[-1].input_channels, layer_scale, 0))
|
||||
self.stage_mask = nn.Sequential(*stage).to(device)
|
||||
|
||||
self.last = nn.Conv2d(
|
||||
block_setting[-1].input_channels,
|
||||
3,
|
||||
kernel_size=3,
|
||||
padding='same',
|
||||
).to(device)
|
||||
|
||||
self.last_mask = 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.stage_smooth(input)
|
||||
x = self.first(x)
|
||||
f = self.layers(x)
|
||||
x = self.stage_pose(f)
|
||||
pose = torch.sigmoid(self.last(x))
|
||||
|
||||
x = self.stage_mask(f)
|
||||
mask = torch.sigmoid(self.last_mask(x))
|
||||
pose = pose * (mask > 0.5).to(torch.float32)
|
||||
|
||||
return pose, mask
|
||||
|
||||
def forward(self, x) -> Tensor:
|
||||
return self._forward_impl(x)
|
||||
|
||||
|
||||
class PoseAE(nn.Module):
|
||||
def __init__(self, features=256, device='cuda'):
|
||||
super().__init__()
|
||||
|
||||
|
||||
block_setting = [
|
||||
CNBlockConfig(96, 192, 3),
|
||||
CNBlockConfig(192, 384, 3),
|
||||
CNBlockConfig(384, features, 9),
|
||||
CNBlockConfig(features, None, 3),
|
||||
]
|
||||
|
||||
block_setting_reverse = [
|
||||
CNBlockConfig(features, 384, 3),
|
||||
CNBlockConfig(384, 192, 9),
|
||||
CNBlockConfig(192, 96, 3),
|
||||
CNBlockConfig(96, 27, 3),
|
||||
CNBlockConfig(27, 9, 3),
|
||||
CNBlockConfig(9, None, 3),
|
||||
]
|
||||
|
||||
self.encoder = ConvNeXt(block_setting, input_features=4, device=device)
|
||||
self.decoder = ConvNeXtReversePose(block_setting_reverse, device=device)
|
||||
|
||||
self.flatten = nn.Flatten()
|
||||
|
||||
|
||||
def forward(self, input) -> torch.Tensor:
|
||||
mask, pose_in = input
|
||||
x, _ = self.encoder(torch.cat([mask, pose_in], dim=1))
|
||||
|
||||
pose_out, mask_out = self.decoder(x)
|
||||
|
||||
return pose_out, mask_out
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
# test the model
|
||||
model = PoseAE()
|
||||
# model = torch.compile(model)
|
||||
print(model)
|
||||
print(sum(p.numel() for p in model.parameters() if p.requires_grad))
|
||||
print(sum(p.numel() for p in model.parameters()))
|
||||
|
||||
x = torch.rand((4, 1, 512, 512)).to('cuda')
|
||||
y = torch.rand((4, 3, 512, 512)).to('cuda')
|
||||
output = model((x, y))
|
||||
|
||||
|
||||
|
||||
pass
|
||||
Reference in New Issue
Block a user