205 lines
6.9 KiB
Python
205 lines
6.9 KiB
Python
import os
|
|
import sys
|
|
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = '7'
|
|
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
|
|
|
|
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 ConvNeXtReverseMaskPose(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:
|
|
# 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)
|
|
|
|
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_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.first(input)
|
|
f = self.layers(x)
|
|
x = self.stage_pose(f)
|
|
pose = self.last(x)
|
|
|
|
x = self.stage_mask(f)
|
|
mask = self.last_mask(x)
|
|
|
|
return pose, mask
|
|
|
|
def forward(self, x) -> Tensor:
|
|
return self._forward_impl(x)
|
|
|
|
|
|
class PoseEncoder(nn.Module):
|
|
def __init__(self, features=128, 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),
|
|
]
|
|
|
|
rotation_block_setting = [
|
|
CNBlockConfig(64, 128, 3),
|
|
CNBlockConfig(128, 128, 3),
|
|
CNBlockConfig(128, 256, 9),
|
|
CNBlockConfig(256, None, 3),
|
|
]
|
|
|
|
rotation_block_setting_reverse = [
|
|
CNBlockConfig(256, 128, 3),
|
|
CNBlockConfig(128, 64, 9),
|
|
CNBlockConfig(64, 32, 3),
|
|
CNBlockConfig(32, 16, 3),
|
|
CNBlockConfig(16, None, 3),
|
|
]
|
|
|
|
self.encoder = ConvNeXt(block_setting, input_features=4, device=device)
|
|
self.decoder = ConvNeXtReverseMaskPose(block_setting_reverse, device=device)
|
|
self.rotation_encoder = ConvNeXt(rotation_block_setting, input_features=7, device=device)
|
|
self.rotation_decoder = ConvNeXt(rotation_block_setting_reverse, reversed=True, device=device)
|
|
|
|
self.flatten = nn.Flatten()
|
|
|
|
|
|
def forward(self, input) -> torch.Tensor:
|
|
mask, pose_in, pose_out = input
|
|
features, _ = self.encoder(torch.cat([mask, pose_in], dim=1))
|
|
x = features.view(features.size(0), features.size(1), -1)
|
|
y, _ = self.rotation_encoder(torch.cat([mask, pose_in, pose_out], dim=1))
|
|
rotation = self.rotation_decoder(y)
|
|
rotation = torch.squeeze(rotation, dim=1)
|
|
rotation = torch.softmax(rotation, dim=1)
|
|
|
|
x = torch.einsum('bij,bkj->bik', x, rotation)
|
|
x = x.view(x.size(0), x.size(1), 16, 16)
|
|
pose, mask = self.decoder(x)
|
|
|
|
return pose, mask, rotation
|
|
|
|
|
|
if __name__ == '__main__':
|
|
# test the model
|
|
model = PoseEncoder()
|
|
# 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')
|
|
z = torch.rand((4, 3, 512, 512)).to('cuda')
|
|
output = model((x, y, z))
|
|
|
|
from torchviz import make_dot
|
|
graph = make_dot(output, params=dict(model.named_parameters()))
|
|
graph.render("model_graph", format="pdf")
|
|
|
|
pass |