fhw
1 year ago
18 changed files with 3317 additions and 2 deletions
@ -0,0 +1,177 @@ |
|||
_base_ = [ |
|||
'../../../../_base_/default_runtime.py', |
|||
'../../../../_base_/datasets/coco.py' |
|||
] |
|||
evaluation = dict(interval=1, metric='mAP', save_best='AP') |
|||
|
|||
optimizer = dict(type='AdamW', |
|||
lr=5e-4, betas=(0.9, 0.999), weight_decay=0.1, |
|||
constructor='LayerDecayOptimizerConstructor', |
|||
paramwise_cfg=dict( |
|||
num_layers=12, |
|||
layer_decay_rate=0.75, |
|||
custom_keys={ |
|||
'bias': dict(decay_multi=0.), |
|||
'pos_embed': dict(decay_mult=0.), |
|||
'relative_position_bias_table': dict(decay_mult=0.), |
|||
'norm': dict(decay_mult=0.) |
|||
} |
|||
) |
|||
) |
|||
|
|||
optimizer_config = dict(grad_clip=dict(max_norm=1., norm_type=2)) |
|||
|
|||
# learning policy |
|||
lr_config = dict( |
|||
policy='step', |
|||
warmup='linear', |
|||
warmup_iters=500, |
|||
warmup_ratio=0.001, |
|||
step=[170, 200]) |
|||
total_epochs = 210 |
|||
target_type = 'GaussianHeatmap' |
|||
channel_cfg = dict( |
|||
num_output_channels=17, |
|||
dataset_joints=17, |
|||
dataset_channel=[ |
|||
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16], |
|||
], |
|||
inference_channel=[ |
|||
0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16 |
|||
]) |
|||
|
|||
# model settings |
|||
model = dict( |
|||
type='TopDownSelf', |
|||
pretrained=None, |
|||
backbone=dict( |
|||
type='ViTSam', |
|||
img_size=(256, 192), |
|||
patch_size=16, |
|||
embed_dim=768, |
|||
depth=12, |
|||
num_heads=12, |
|||
ratio=1, |
|||
use_checkpoint=False, |
|||
mlp_ratio=4, |
|||
qkv_bias=True, |
|||
drop_path_rate=0.3, |
|||
frozen_stages=12, |
|||
freeze_attn = True, |
|||
freeze_ffn = True, |
|||
samvit_checkpoint='/root/autodl-tmp/code/ViTPose/checkpoints/sam/sam_vit_b_01ec64.pth' |
|||
), |
|||
keypoint_head=dict( |
|||
type='TopdownHeatmapSimpleHead', |
|||
in_channels=768, |
|||
num_deconv_layers=2, |
|||
num_deconv_filters=(256, 256), |
|||
num_deconv_kernels=(4, 4), |
|||
extra=dict(final_conv_kernel=1, ), |
|||
out_channels=channel_cfg['num_output_channels'], |
|||
loss_keypoint=dict(type='JointsMSELoss', use_target_weight=True)), |
|||
train_cfg=dict(), |
|||
test_cfg=dict( |
|||
flip_test=True, |
|||
post_process='default', |
|||
shift_heatmap=False, |
|||
target_type=target_type, |
|||
modulate_kernel=11, |
|||
use_udp=True)) |
|||
|
|||
data_root = '/root/autodl-tmp/dataset/coco2017/' |
|||
|
|||
data_cfg = dict( |
|||
image_size=[192, 256], |
|||
heatmap_size=[48, 64], |
|||
num_output_channels=channel_cfg['num_output_channels'], |
|||
num_joints=channel_cfg['dataset_joints'], |
|||
dataset_channel=channel_cfg['dataset_channel'], |
|||
inference_channel=channel_cfg['inference_channel'], |
|||
soft_nms=False, |
|||
nms_thr=1.0, |
|||
oks_thr=0.9, |
|||
vis_thr=0.2, |
|||
use_gt_bbox=False, |
|||
det_bbox_thr=0.0, |
|||
bbox_file=f'{data_root}/person_detection_results/COCO_val2017_detections_AP_H_56_person.json', |
|||
) |
|||
|
|||
train_pipeline = [ |
|||
dict(type='LoadImageFromFile'), |
|||
dict(type='TopDownRandomFlip', flip_prob=0.5), |
|||
dict( |
|||
type='TopDownHalfBodyTransform', |
|||
num_joints_half_body=8, |
|||
prob_half_body=0.3), |
|||
dict( |
|||
type='TopDownGetRandomScaleRotation', rot_factor=40, scale_factor=0.5), |
|||
# dict(type='TopDownAffine', use_udp=True), |
|||
dict(type='TopDownAffineSam', use_udp=True), |
|||
dict(type='ToTensorSam'), |
|||
dict( |
|||
type='NormalizeTensorSam', |
|||
mean=[0.485, 0.456, 0.406], |
|||
std=[0.229, 0.224, 0.225]), |
|||
dict( |
|||
type='TopDownGenerateTarget', |
|||
sigma=2, |
|||
encoding='UDP', |
|||
target_type=target_type), |
|||
dict( |
|||
type='Collect', |
|||
keys=['img', 'sam_img', 'target', 'target_weight'], |
|||
meta_keys=[ |
|||
'image_file', 'joints_3d', 'joints_3d_visible', 'center', 'scale', |
|||
'rotation', 'bbox_score', 'flip_pairs' |
|||
]), |
|||
] |
|||
|
|||
val_pipeline = [ |
|||
dict(type='LoadImageFromFile'), |
|||
# dict(type='TopDownAffine', use_udp=True), |
|||
dict(type='TopDownAffineSam', use_udp=True), |
|||
dict(type='ToTensorSam'), |
|||
dict( |
|||
type='NormalizeTensorSam', |
|||
mean=[0.485, 0.456, 0.406], |
|||
std=[0.229, 0.224, 0.225]), |
|||
dict( |
|||
type='Collect', |
|||
keys=['img', 'sam_img'], |
|||
meta_keys=[ |
|||
'image_file', 'center', 'scale', 'rotation', 'bbox_score', |
|||
'flip_pairs' |
|||
]), |
|||
] |
|||
|
|||
test_pipeline = val_pipeline |
|||
|
|||
data = dict( |
|||
samples_per_gpu=12, |
|||
workers_per_gpu=4, |
|||
val_dataloader=dict(samples_per_gpu=12), |
|||
test_dataloader=dict(samples_per_gpu=12), |
|||
train=dict( |
|||
type='TopDownCocoDataset', |
|||
ann_file=f'{data_root}/annotations/person_keypoints_train2017.json', |
|||
img_prefix=f'{data_root}/train2017/', |
|||
data_cfg=data_cfg, |
|||
pipeline=train_pipeline, |
|||
dataset_info={{_base_.dataset_info}}), |
|||
val=dict( |
|||
type='TopDownCocoDataset', |
|||
ann_file=f'{data_root}/annotations/person_keypoints_val2017.json', |
|||
img_prefix=f'{data_root}/val2017/', |
|||
data_cfg=data_cfg, |
|||
pipeline=val_pipeline, |
|||
dataset_info={{_base_.dataset_info}}), |
|||
test=dict( |
|||
type='TopDownCocoDataset', |
|||
ann_file=f'{data_root}/annotations/person_keypoints_val2017.json', |
|||
img_prefix=f'{data_root}/val2017/', |
|||
data_cfg=data_cfg, |
|||
pipeline=test_pipeline, |
|||
dataset_info={{_base_.dataset_info}}), |
|||
) |
|||
|
|||
@ -0,0 +1 @@ |
|||
../../configs |
|||
@ -0,0 +1 @@ |
|||
../../demo |
|||
@ -0,0 +1 @@ |
|||
../../model-index.yml |
|||
@ -0,0 +1 @@ |
|||
../../tools |
|||
@ -0,0 +1,76 @@ |
|||
# Copyright (c) OpenMMLab. All rights reserved. |
|||
import warnings |
|||
from collections.abc import Sequence |
|||
|
|||
import mmcv |
|||
import numpy as np |
|||
from mmcv.parallel import DataContainer as DC |
|||
from mmcv.utils import build_from_cfg |
|||
from numpy import random |
|||
from torchvision.transforms import functional as F |
|||
|
|||
from ..builder import PIPELINES |
|||
|
|||
try: |
|||
import albumentations |
|||
except ImportError: |
|||
albumentations = None |
|||
|
|||
|
|||
@PIPELINES.register_module() |
|||
class ToTensorSam: |
|||
"""Transform image to Tensor. |
|||
|
|||
Required key: 'img'. Modifies key: 'img'. |
|||
|
|||
Args: |
|||
results (dict): contain all information about training. |
|||
""" |
|||
|
|||
def __call__(self, results): |
|||
if isinstance(results['img'], (list, tuple)): |
|||
results['img'] = [F.to_tensor(img) for img in results['img']] |
|||
# 修改 |
|||
results['sam_img'] = [F.to_tensor(sam_img) for sam_img in results['sam_img']] |
|||
else: |
|||
results['img'] = F.to_tensor(results['img']) |
|||
# 修改 |
|||
results['sam_img'] = F.to_tensor(results['sam_img']) |
|||
|
|||
return results |
|||
|
|||
|
|||
@PIPELINES.register_module() |
|||
class NormalizeTensorSam: |
|||
"""Normalize the Tensor image (CxHxW), with mean and std. |
|||
|
|||
Required key: 'img'. Modifies key: 'img'. |
|||
|
|||
Args: |
|||
mean (list[float]): Mean values of 3 channels. |
|||
std (list[float]): Std values of 3 channels. |
|||
""" |
|||
|
|||
def __init__(self, mean, std): |
|||
self.mean = mean |
|||
self.std = std |
|||
|
|||
def __call__(self, results): |
|||
if isinstance(results['img'], (list, tuple)): |
|||
results['img'] = [ |
|||
F.normalize(img, mean=self.mean, std=self.std) |
|||
for img in results['img'] |
|||
] |
|||
# 修改 |
|||
results['sam_img'] = [ |
|||
F.normalize(sam_img, mean=self.mean, std=self.std) |
|||
for sam_img in results['sam_img'] |
|||
] |
|||
else: |
|||
results['img'] = F.normalize( |
|||
results['img'], mean=self.mean, std=self.std) |
|||
# 修改 |
|||
results['sam_img'] = F.normalize( |
|||
results['sam_img'], mean=self.mean, std=self.std) |
|||
|
|||
return results |
|||
@ -0,0 +1,113 @@ |
|||
import cv2 |
|||
import numpy as np |
|||
|
|||
from mmpose.core.post_processing import (affine_transform, fliplr_joints, |
|||
get_affine_transform, get_warp_matrix, |
|||
warp_affine_joints) |
|||
from mmpose.datasets.builder import PIPELINES |
|||
|
|||
@PIPELINES.register_module() |
|||
class TopDownAffineSam: |
|||
"""Affine transform the image to make input. |
|||
|
|||
Required keys:'img', 'joints_3d', 'joints_3d_visible', 'ann_info','scale', |
|||
'rotation' and 'center'. |
|||
|
|||
Modified keys:'img', 'joints_3d', and 'joints_3d_visible'. |
|||
|
|||
Args: |
|||
use_udp (bool): To use unbiased data processing. |
|||
Paper ref: Huang et al. The Devil is in the Details: Delving into |
|||
Unbiased Data Processing for Human Pose Estimation (CVPR 2020). |
|||
""" |
|||
|
|||
def __init__(self, use_udp=False): |
|||
self.use_udp = use_udp |
|||
|
|||
def __call__(self, results): |
|||
image_size = results['ann_info']['image_size'] |
|||
# 修改 |
|||
sam_image_size = np.array([1024, 1024]) |
|||
|
|||
img = results['img'] |
|||
joints_3d = results['joints_3d'] |
|||
joints_3d_visible = results['joints_3d_visible'] |
|||
c = results['center'] |
|||
s = results['scale'] |
|||
r = results['rotation'] |
|||
# 修改 |
|||
sam_img = img |
|||
|
|||
if self.use_udp: |
|||
trans = get_warp_matrix(r, c * 2.0, image_size - 1.0, s * 200.0) |
|||
# 修改 |
|||
sam_trans = get_warp_matrix(r, c * 2.0, sam_image_size - 1.0, s * 200.0) |
|||
if not isinstance(img, list): |
|||
img = cv2.warpAffine( |
|||
img, |
|||
trans, (int(image_size[0]), int(image_size[1])), |
|||
flags=cv2.INTER_LINEAR) |
|||
# 修改 |
|||
sam_img = cv2.warpAffine( |
|||
sam_img, |
|||
sam_trans, (int(sam_image_size[0]), int(sam_image_size[1])), |
|||
flags=cv2.INTER_LINEAR) |
|||
else: |
|||
img = [ |
|||
cv2.warpAffine( |
|||
i, |
|||
trans, (int(image_size[0]), int(image_size[1])), |
|||
flags=cv2.INTER_LINEAR) for i in img |
|||
] |
|||
# 修改 |
|||
sam_img = [ |
|||
cv2.warpAffine( |
|||
i, |
|||
sam_trans, (int(sam_image_size[0]), int(sam_image_size[1])), |
|||
flags=cv2.INTER_LINEAR) for i in sam_img |
|||
] |
|||
|
|||
joints_3d[:, 0:2] = \ |
|||
warp_affine_joints(joints_3d[:, 0:2].copy(), trans) |
|||
|
|||
else: |
|||
trans = get_affine_transform(c, s, r, image_size) |
|||
# 修改 |
|||
sam_trans = get_affine_transform(c, s, r, sam_image_size) |
|||
if not isinstance(img, list): |
|||
img = cv2.warpAffine( |
|||
img, |
|||
trans, (int(image_size[0]), int(image_size[1])), |
|||
flags=cv2.INTER_LINEAR) |
|||
|
|||
# 修改 |
|||
sam_img = cv2.warpAffine( |
|||
sam_img, |
|||
sam_trans, (int(sam_image_size[0]), int(sam_image_size[1])), |
|||
flags=cv2.INTER_LINEAR) |
|||
else: |
|||
img = [ |
|||
cv2.warpAffine( |
|||
i, |
|||
trans, (int(image_size[0]), int(image_size[1])), |
|||
flags=cv2.INTER_LINEAR) for i in img |
|||
] |
|||
# 修改 |
|||
sam_img = [ |
|||
cv2.warpAffine( |
|||
i, |
|||
sam_trans, (int(sam_image_size[0]), int(sam_image_size[1])), |
|||
flags=cv2.INTER_LINEAR) for i in sam_img |
|||
] |
|||
|
|||
for i in range(results['ann_info']['num_joints']): |
|||
if joints_3d_visible[i, 0] > 0.0: |
|||
joints_3d[i, |
|||
0:2] = affine_transform(joints_3d[i, 0:2], trans) |
|||
|
|||
results['img'] = img |
|||
results['sam_img'] = sam_img |
|||
results['joints_3d'] = joints_3d |
|||
results['joints_3d_visible'] = joints_3d_visible |
|||
|
|||
return results |
|||
@ -0,0 +1 @@ |
|||
from .image_encoder import build_vit_sam |
|||
@ -0,0 +1,477 @@ |
|||
# -------------------------------------------------------------------- |
|||
# Copyright (c) Meta Platforms, Inc. and affiliates. |
|||
# All rights reserved. |
|||
|
|||
# This source code is licensed under the license found in the |
|||
# LICENSE file in the root directory of this source tree. |
|||
# -------------------------------------------------------------------- |
|||
|
|||
from typing import Optional, Tuple, Type |
|||
|
|||
import torch |
|||
import torch.nn as nn |
|||
import torch.nn |
|||
import torch.nn.functional as F |
|||
|
|||
from functools import partial |
|||
|
|||
|
|||
# ---------------------- Vision Transformer of Segment-Anything ---------------------- |
|||
class ImageEncoderViT(nn.Module): |
|||
""" |
|||
We remove the neck which used in the Segment-Anything. |
|||
""" |
|||
def __init__(self, |
|||
img_size : int = 1024, |
|||
patch_size : int = 16, |
|||
in_chans : int = 3, |
|||
embed_dim : int = 768, |
|||
depth : int = 12, |
|||
num_heads : int = 12, |
|||
mlp_ratio : float = 4.0, |
|||
qkv_bias : bool = True, |
|||
norm_layer : Type[nn.Module] = nn.LayerNorm, |
|||
act_layer : Type[nn.Module] = nn.GELU, |
|||
use_abs_pos : bool = True, |
|||
use_rel_pos : bool = True, |
|||
window_size : int = 0, |
|||
global_attn_indexes : Tuple[int, ...] = (), |
|||
checkpoint = None |
|||
) -> None: |
|||
super().__init__() |
|||
self.img_size = img_size |
|||
self.patch_size = patch_size |
|||
self.embed_dim = embed_dim |
|||
self.num_patches = (img_size // patch_size) ** 2 |
|||
# self.num_patches = (img_size[0] // patch_size) * (img_size[1] // patch_size) |
|||
self.pos_embed: Optional[nn.Parameter] = None |
|||
self.checkpoint = checkpoint |
|||
if use_abs_pos: |
|||
# Initialize absolute positional embedding with pretrain image size. |
|||
self.pos_embed = nn.Parameter( |
|||
torch.zeros(1, img_size // patch_size, img_size // patch_size, embed_dim) |
|||
) |
|||
|
|||
# ------------ Model parameters ------------ |
|||
## Patch embedding layer |
|||
self.patch_embed = PatchEmbed( |
|||
kernel_size=(patch_size, patch_size), |
|||
stride=(patch_size, patch_size), |
|||
in_chans=in_chans, |
|||
embed_dim=embed_dim, |
|||
) |
|||
|
|||
## ViT blocks |
|||
self.blocks = nn.ModuleList() |
|||
for i in range(depth): |
|||
block = Block(dim = embed_dim, |
|||
num_heads = num_heads, |
|||
mlp_ratio = mlp_ratio, |
|||
qkv_bias = qkv_bias, |
|||
norm_layer = norm_layer, |
|||
act_layer = act_layer, |
|||
use_rel_pos = use_rel_pos, |
|||
window_size = window_size if i not in global_attn_indexes else 0, |
|||
input_size = (img_size // patch_size, img_size // patch_size), |
|||
) |
|||
self.blocks.append(block) |
|||
|
|||
self.load_pretrained() |
|||
|
|||
def load_pretrained(self): |
|||
if self.checkpoint is not None: |
|||
print('Loading SAM pretrained weight from : {}'.format(self.checkpoint)) |
|||
# checkpoint state dict |
|||
checkpoint_state_dict = torch.load(self.checkpoint, map_location="cpu") |
|||
# model state dict |
|||
model_state_dict = self.state_dict() |
|||
encoder_state_dict = {} |
|||
# check |
|||
for k in list(checkpoint_state_dict.keys()): |
|||
if "image_encoder" in k and k[14:] in model_state_dict: |
|||
shape_model = tuple(model_state_dict[k[14:]].shape) |
|||
shape_checkpoint = tuple(checkpoint_state_dict[k].shape) |
|||
if shape_model == shape_checkpoint or "pos_embed" in k: |
|||
encoder_state_dict[k[14:]] = checkpoint_state_dict[k] |
|||
else: |
|||
print("Shape unmatch: ", k) |
|||
|
|||
# interpolate position embedding |
|||
# interpolate_pos_embed(self, encoder_state_dict, ((self.img_size[0] // self.patch_size), (self.img_size[1] // self.patch_size))) |
|||
interpolate_pos_embed(self, encoder_state_dict,) |
|||
|
|||
# load the weight |
|||
self.load_state_dict(encoder_state_dict, strict=False) |
|||
else: |
|||
print('No SAM pretrained.') |
|||
|
|||
# @torch.no_grad() |
|||
def forward(self, x: torch.Tensor) -> torch.Tensor: |
|||
# with torch.no_grad(): |
|||
|
|||
x = self.patch_embed(x) |
|||
if self.pos_embed is not None: |
|||
x = x + self.pos_embed |
|||
|
|||
for blk in self.blocks: |
|||
x = blk(x) |
|||
|
|||
# [B, H, W, C] -> [B, N, C] |
|||
return x.flatten(1, 2) |
|||
|
|||
|
|||
# ---------------------- Model modules ---------------------- |
|||
class MLPBlock(nn.Module): |
|||
def __init__(self, |
|||
embedding_dim: int, |
|||
mlp_dim: int, |
|||
act: Type[nn.Module] = nn.GELU, |
|||
) -> None: |
|||
super().__init__() |
|||
self.lin1 = nn.Linear(embedding_dim, mlp_dim) |
|||
self.lin2 = nn.Linear(mlp_dim, embedding_dim) |
|||
self.act = act() |
|||
|
|||
def forward(self, x: torch.Tensor) -> torch.Tensor: |
|||
return self.lin2(self.act(self.lin1(x))) |
|||
|
|||
class LayerNorm2d(nn.Module): |
|||
def __init__(self, num_channels: int, eps: float = 1e-6) -> None: |
|||
super().__init__() |
|||
self.weight = nn.Parameter(torch.ones(num_channels)) |
|||
self.bias = nn.Parameter(torch.zeros(num_channels)) |
|||
self.eps = eps |
|||
|
|||
def forward(self, x: torch.Tensor) -> torch.Tensor: |
|||
u = x.mean(1, keepdim=True) |
|||
s = (x - u).pow(2).mean(1, keepdim=True) |
|||
x = (x - u) / torch.sqrt(s + self.eps) |
|||
x = self.weight[:, None, None] * x + self.bias[:, None, None] |
|||
|
|||
return x |
|||
|
|||
class Block(nn.Module): |
|||
def __init__(self, |
|||
dim : int, |
|||
num_heads : int, |
|||
mlp_ratio : float = 4.0, |
|||
qkv_bias : bool = True, |
|||
norm_layer : Type[nn.Module] = nn.LayerNorm, |
|||
act_layer : Type[nn.Module] = nn.GELU, |
|||
use_rel_pos : bool = False, |
|||
window_size : int = 0, |
|||
input_size : Optional[Tuple[int, int]] = None, |
|||
) -> None: |
|||
super().__init__() |
|||
# -------------- Basic parameters -------------- |
|||
self.window_size = window_size |
|||
# -------------- Model parameters -------------- |
|||
self.norm1 = norm_layer(dim) |
|||
self.attn = Attention(dim = dim, |
|||
num_heads = num_heads, |
|||
qkv_bias = qkv_bias, |
|||
use_rel_pos = use_rel_pos, |
|||
input_size = input_size if window_size == 0 else (window_size, window_size), |
|||
) |
|||
self.norm2 = norm_layer(dim) |
|||
self.mlp = MLPBlock(embedding_dim=dim, mlp_dim=int(dim * mlp_ratio), act=act_layer) |
|||
|
|||
def forward(self, x: torch.Tensor) -> torch.Tensor: |
|||
shortcut = x |
|||
x = self.norm1(x) |
|||
# Window partition |
|||
if self.window_size > 0: |
|||
H, W = x.shape[1], x.shape[2] |
|||
x, pad_hw = window_partition(x, self.window_size) |
|||
|
|||
x = self.attn(x) |
|||
# Reverse window partition |
|||
if self.window_size > 0: |
|||
x = window_unpartition(x, self.window_size, pad_hw, (H, W)) |
|||
|
|||
x = shortcut + x |
|||
x = x + self.mlp(self.norm2(x)) |
|||
|
|||
return x |
|||
|
|||
class Attention(nn.Module): |
|||
def __init__(self, |
|||
dim: int, |
|||
num_heads: int = 8, |
|||
qkv_bias: bool = True, |
|||
use_rel_pos: bool = False, |
|||
input_size: Optional[Tuple[int, int]] = None, |
|||
) -> None: |
|||
super().__init__() |
|||
# -------------- Basic parameters -------------- |
|||
self.num_heads = num_heads |
|||
head_dim = dim // num_heads |
|||
self.scale = head_dim**-0.5 |
|||
self.use_rel_pos = use_rel_pos |
|||
if self.use_rel_pos: |
|||
assert ( |
|||
input_size is not None |
|||
), "Input size must be provided if using relative positional encoding." |
|||
# initialize relative positional embeddings |
|||
self.rel_pos_h = nn.Parameter(torch.zeros(2 * input_size[0] - 1, head_dim)) |
|||
self.rel_pos_w = nn.Parameter(torch.zeros(2 * input_size[1] - 1, head_dim)) |
|||
|
|||
# -------------- Model parameters -------------- |
|||
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) |
|||
self.proj = nn.Linear(dim, dim) |
|||
|
|||
def forward(self, x: torch.Tensor) -> torch.Tensor: |
|||
B, H, W, _ = x.shape |
|||
# qkv with shape (3, B, nHead, H * W, C) |
|||
qkv = self.qkv(x).reshape(B, H * W, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4) |
|||
# q, k, v with shape (B * nHead, H * W, C) |
|||
q, k, v = qkv.reshape(3, B * self.num_heads, H * W, -1).unbind(0) |
|||
|
|||
attn = (q * self.scale) @ k.transpose(-2, -1) |
|||
|
|||
if self.use_rel_pos: |
|||
attn = add_decomposed_rel_pos(attn, q, self.rel_pos_h, self.rel_pos_w, (H, W), (H, W)) |
|||
|
|||
attn = attn.softmax(dim=-1) |
|||
x = (attn @ v).view(B, self.num_heads, H, W, -1).permute(0, 2, 3, 1, 4).reshape(B, H, W, -1) |
|||
x = self.proj(x) |
|||
|
|||
return x |
|||
|
|||
class PatchEmbed(nn.Module): |
|||
def __init__(self, |
|||
kernel_size : Tuple[int, int] = (16, 16), |
|||
stride : Tuple[int, int] = (16, 16), |
|||
padding : Tuple[int, int] = (0, 0), |
|||
in_chans : int = 3, |
|||
embed_dim : int = 768, |
|||
) -> None: |
|||
super().__init__() |
|||
self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=kernel_size, stride=stride, padding=padding) |
|||
|
|||
def forward(self, x: torch.Tensor) -> torch.Tensor: |
|||
x = self.proj(x) |
|||
# [B, C, H, W] -> [B, H, W, C] |
|||
x = x.permute(0, 2, 3, 1) |
|||
|
|||
return x |
|||
|
|||
|
|||
# ---------------------- Model functions ---------------------- |
|||
def window_partition(x: torch.Tensor, |
|||
window_size: int, |
|||
) -> Tuple[torch.Tensor, Tuple[int, int]]: |
|||
""" |
|||
Partition into non-overlapping windows with padding if needed. |
|||
Args: |
|||
x (tensor): input tokens with [B, H, W, C]. |
|||
window_size (int): window size. |
|||
|
|||
Returns: |
|||
windows: windows after partition with [B * num_windows, window_size, window_size, C]. |
|||
(Hp, Wp): padded height and width before partition |
|||
""" |
|||
B, H, W, C = x.shape |
|||
|
|||
pad_h = (window_size - H % window_size) % window_size |
|||
pad_w = (window_size - W % window_size) % window_size |
|||
if pad_h > 0 or pad_w > 0: |
|||
x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h)) |
|||
Hp, Wp = H + pad_h, W + pad_w |
|||
|
|||
x = x.view(B, Hp // window_size, window_size, Wp // window_size, window_size, C) |
|||
windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C) |
|||
|
|||
return windows, (Hp, Wp) |
|||
|
|||
def window_unpartition(windows: torch.Tensor, |
|||
window_size: int, |
|||
pad_hw: Tuple[int, int], |
|||
hw: Tuple[int, int], |
|||
) -> torch.Tensor: |
|||
""" |
|||
Window unpartition into original sequences and removing padding. |
|||
Args: |
|||
windows (tensor): input tokens with [B * num_windows, window_size, window_size, C]. |
|||
window_size (int): window size. |
|||
pad_hw (Tuple): padded height and width (Hp, Wp). |
|||
hw (Tuple): original height and width (H, W) before padding. |
|||
|
|||
Returns: |
|||
x: unpartitioned sequences with [B, H, W, C]. |
|||
""" |
|||
Hp, Wp = pad_hw |
|||
H, W = hw |
|||
B = windows.shape[0] // (Hp * Wp // window_size // window_size) |
|||
x = windows.view(B, Hp // window_size, Wp // window_size, window_size, window_size, -1) |
|||
x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, Hp, Wp, -1) |
|||
|
|||
if Hp > H or Wp > W: |
|||
x = x[:, :H, :W, :].contiguous() |
|||
|
|||
return x |
|||
|
|||
def get_rel_pos(q_size: int, |
|||
k_size: int, |
|||
rel_pos: torch.Tensor, |
|||
)-> torch.Tensor: |
|||
""" |
|||
Get relative positional embeddings according to the relative positions of |
|||
query and key sizes. |
|||
Args: |
|||
q_size (int): size of query q. |
|||
k_size (int): size of key k. |
|||
rel_pos (Tensor): relative position embeddings (L, C). |
|||
|
|||
Returns: |
|||
Extracted positional embeddings according to relative positions. |
|||
""" |
|||
max_rel_dist = int(2 * max(q_size, k_size) - 1) |
|||
# Interpolate rel pos if needed. |
|||
if rel_pos.shape[0] != max_rel_dist: |
|||
# Interpolate rel pos. |
|||
rel_pos_resized = F.interpolate( |
|||
rel_pos.reshape(1, rel_pos.shape[0], -1).permute(0, 2, 1), |
|||
size=max_rel_dist, |
|||
mode="linear", |
|||
) |
|||
rel_pos_resized = rel_pos_resized.reshape(-1, max_rel_dist).permute(1, 0) |
|||
else: |
|||
rel_pos_resized = rel_pos |
|||
|
|||
# Scale the coords with short length if shapes for q and k are different. |
|||
q_coords = torch.arange(q_size)[:, None] * max(k_size / q_size, 1.0) |
|||
k_coords = torch.arange(k_size)[None, :] * max(q_size / k_size, 1.0) |
|||
relative_coords = (q_coords - k_coords) + (k_size - 1) * max(q_size / k_size, 1.0) |
|||
|
|||
return rel_pos_resized[relative_coords.long()] |
|||
|
|||
def add_decomposed_rel_pos(attn : torch.Tensor, |
|||
q : torch.Tensor, |
|||
rel_pos_h : torch.Tensor, |
|||
rel_pos_w : torch.Tensor, |
|||
q_size : Tuple[int, int], |
|||
k_size : Tuple[int, int], |
|||
) -> torch.Tensor: |
|||
q_h, q_w = q_size |
|||
k_h, k_w = k_size |
|||
Rh = get_rel_pos(q_h, k_h, rel_pos_h) |
|||
Rw = get_rel_pos(q_w, k_w, rel_pos_w) |
|||
|
|||
B, _, dim = q.shape |
|||
r_q = q.reshape(B, q_h, q_w, dim) |
|||
rel_h = torch.einsum("bhwc,hkc->bhwk", r_q, Rh) |
|||
rel_w = torch.einsum("bhwc,wkc->bhwk", r_q, Rw) |
|||
|
|||
attn = ( |
|||
attn.view(B, q_h, q_w, k_h, k_w) + rel_h[:, :, :, :, None] + rel_w[:, :, :, None, :] |
|||
).view(B, q_h * q_w, k_h * k_w) |
|||
|
|||
return attn |
|||
|
|||
def interpolate_pos_embed(model, checkpoint_model): |
|||
if 'pos_embed' in checkpoint_model: |
|||
# Pos embed from checkpoint |
|||
pos_embed_checkpoint = checkpoint_model['pos_embed'] |
|||
embedding_size = pos_embed_checkpoint.shape[-1] |
|||
# Pos embed from model |
|||
pos_embed_model = model.pos_embed |
|||
num_patches = model.num_patches |
|||
# [B, H, W, C] -> [B, N, C] |
|||
pos_embed_checkpoint = pos_embed_checkpoint.flatten(1, 2) |
|||
pos_embed_model = pos_embed_model.flatten(1, 2) |
|||
|
|||
orig_num_postions = pos_embed_model.shape[-2] |
|||
num_extra_tokens = orig_num_postions - num_patches |
|||
|
|||
# height (== width) for the checkpoint position embedding |
|||
orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5) |
|||
new_size = int(num_patches ** 0.5) |
|||
|
|||
# height (== width) for the new position embedding |
|||
# class_token and dist_token are kept unchanged |
|||
if orig_size != new_size: |
|||
print("- Position interpolate from %dx%d to %dx%d" % (orig_size, orig_size, new_size, new_size)) |
|||
extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens] |
|||
# only the position tokens are interpolated |
|||
pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:] |
|||
pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2) |
|||
pos_tokens = torch.nn.functional.interpolate(pos_tokens, |
|||
size=(new_size,new_size), |
|||
mode='bicubic', |
|||
align_corners=False) |
|||
pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2) |
|||
new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1) |
|||
new_pos_embed = new_pos_embed.reshape(-1, int(orig_num_postions ** 0.5), int(orig_num_postions ** 0.5), embedding_size) |
|||
checkpoint_model['pos_embed'] = new_pos_embed |
|||
|
|||
|
|||
# ------------------------ Model Functions ------------------------ |
|||
def build_vit_sam(model_name="vit_h", img_size=1024, patch_size=16, img_dim=3, checkpoint=None): |
|||
if model_name == "vit_b": |
|||
return ImageEncoderViT(img_size=img_size, |
|||
patch_size=patch_size, |
|||
in_chans=img_dim, |
|||
embed_dim=768, |
|||
depth=12, |
|||
num_heads=12, |
|||
mlp_ratio=4.0, |
|||
norm_layer=partial(nn.LayerNorm, eps=1e-6), |
|||
global_attn_indexes=[2, 5, 8, 11], |
|||
window_size=14, |
|||
checkpoint=checkpoint, |
|||
) |
|||
if model_name == "vit_l": |
|||
return ImageEncoderViT(img_size=img_size, |
|||
patch_size=patch_size, |
|||
in_chans=img_dim, |
|||
embed_dim=1024, |
|||
depth=24, |
|||
num_heads=16, |
|||
mlp_ratio=4.0, |
|||
norm_layer=partial(nn.LayerNorm, eps=1e-6), |
|||
global_attn_indexes=[5, 11, 17, 23], |
|||
window_size=14, |
|||
checkpoint=checkpoint, |
|||
) |
|||
if model_name == "vit_h": |
|||
return ImageEncoderViT(img_size=img_size, |
|||
patch_size=patch_size, |
|||
in_chans=img_dim, |
|||
embed_dim=1280, |
|||
depth=32, |
|||
num_heads=16, |
|||
mlp_ratio=4.0, |
|||
norm_layer=partial(nn.LayerNorm, eps=1e-6), |
|||
global_attn_indexes=[7, 15, 23, 31], |
|||
window_size=14, |
|||
checkpoint=checkpoint, |
|||
) |
|||
|
|||
|
|||
if __name__ == '__main__': |
|||
import torch |
|||
from thop import profile |
|||
|
|||
# Prepare an image as the input |
|||
bs, c, h, w = 2, 3, 1024, 1024 |
|||
x = torch.randn(bs, c, h, w) |
|||
patch_size = 16 |
|||
device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu') |
|||
|
|||
# Build model |
|||
model = build_vit_sam(model_name='vit_b', checkpoint="/home/fhw/code/ViTPose/checkpoints/sam/sam_vit_b_01ec64.pth") |
|||
if torch.cuda.is_available(): |
|||
x = x.to(device) |
|||
model = model.to(device) |
|||
|
|||
# Inference |
|||
outputs = model(x) |
|||
print(outputs.shape) |
|||
|
|||
# Compute FLOPs & Params |
|||
print('==============================') |
|||
model.eval() |
|||
flops, params = profile(model, inputs=(x, ), verbose=False) |
|||
print('GFLOPs : {:.2f}'.format(flops / 1e9 * 2)) |
|||
print('Params : {:.2f} M'.format(params / 1e6)) |
|||
@ -0,0 +1,483 @@ |
|||
# Copyright (c) OpenMMLab. All rights reserved. |
|||
import math |
|||
|
|||
import torch |
|||
from functools import partial |
|||
import torch.nn as nn |
|||
import torch.nn.functional as F |
|||
import torch.utils.checkpoint as checkpoint |
|||
|
|||
from timm.models.layers import drop_path, to_2tuple, trunc_normal_ |
|||
|
|||
from ..builder import BACKBONES |
|||
from .base_backbone import BaseBackbone |
|||
|
|||
from .sam_vit import build_vit_sam |
|||
|
|||
|
|||
def get_abs_pos(abs_pos, h, w, ori_h, ori_w, has_cls_token=True): |
|||
""" |
|||
Calculate absolute positional embeddings. If needed, resize embeddings and remove cls_token |
|||
dimension for the original embeddings. |
|||
Args: |
|||
abs_pos (Tensor): absolute positional embeddings with (1, num_position, C). |
|||
has_cls_token (bool): If true, has 1 embedding in abs_pos for cls token. |
|||
hw (Tuple): size of input image tokens. |
|||
|
|||
Returns: |
|||
Absolute positional embeddings after processing with shape (1, H, W, C) |
|||
""" |
|||
cls_token = None |
|||
B, L, C = abs_pos.shape |
|||
if has_cls_token: |
|||
cls_token = abs_pos[:, 0:1] |
|||
abs_pos = abs_pos[:, 1:] |
|||
|
|||
if ori_h != h or ori_w != w: |
|||
new_abs_pos = F.interpolate( |
|||
abs_pos.reshape(1, ori_h, ori_w, -1).permute(0, 3, 1, 2), |
|||
size=(h, w), |
|||
mode="bicubic", |
|||
align_corners=False, |
|||
).permute(0, 2, 3, 1).reshape(B, -1, C) |
|||
|
|||
else: |
|||
new_abs_pos = abs_pos |
|||
|
|||
if cls_token is not None: |
|||
new_abs_pos = torch.cat([cls_token, new_abs_pos], dim=1) |
|||
return new_abs_pos |
|||
|
|||
class DropPath(nn.Module): |
|||
"""Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). |
|||
""" |
|||
def __init__(self, drop_prob=None): |
|||
super(DropPath, self).__init__() |
|||
self.drop_prob = drop_prob |
|||
|
|||
def forward(self, x): |
|||
return drop_path(x, self.drop_prob, self.training) |
|||
|
|||
def extra_repr(self): |
|||
return 'p={}'.format(self.drop_prob) |
|||
|
|||
class Mlp(nn.Module): |
|||
def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): |
|||
super().__init__() |
|||
out_features = out_features or in_features |
|||
hidden_features = hidden_features or in_features |
|||
self.fc1 = nn.Linear(in_features, hidden_features) |
|||
self.act = act_layer() |
|||
self.fc2 = nn.Linear(hidden_features, out_features) |
|||
self.drop = nn.Dropout(drop) |
|||
|
|||
def forward(self, x): |
|||
x = self.fc1(x) |
|||
x = self.act(x) |
|||
x = self.fc2(x) |
|||
x = self.drop(x) |
|||
return x |
|||
|
|||
class Attention(nn.Module): |
|||
def __init__( |
|||
self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., |
|||
proj_drop=0., attn_head_dim=None,): |
|||
super().__init__() |
|||
self.num_heads = num_heads |
|||
head_dim = dim // num_heads |
|||
self.dim = dim |
|||
|
|||
if attn_head_dim is not None: |
|||
head_dim = attn_head_dim |
|||
all_head_dim = head_dim * self.num_heads |
|||
|
|||
self.scale = qk_scale or head_dim ** -0.5 |
|||
|
|||
self.qkv = nn.Linear(dim, all_head_dim * 3, bias=qkv_bias) |
|||
|
|||
self.attn_drop = nn.Dropout(attn_drop) |
|||
self.proj = nn.Linear(all_head_dim, dim) |
|||
self.proj_drop = nn.Dropout(proj_drop) |
|||
|
|||
def forward(self, x): |
|||
B, N, C = x.shape |
|||
qkv = self.qkv(x) |
|||
qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4) |
|||
q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple) |
|||
|
|||
q = q * self.scale |
|||
attn = (q @ k.transpose(-2, -1)) |
|||
|
|||
attn = attn.softmax(dim=-1) |
|||
attn = self.attn_drop(attn) |
|||
|
|||
x = (attn @ v).transpose(1, 2).reshape(B, N, -1) |
|||
x = self.proj(x) |
|||
x = self.proj_drop(x) |
|||
|
|||
return x |
|||
|
|||
class Block(nn.Module): |
|||
|
|||
def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, |
|||
drop=0., attn_drop=0., drop_path=0., act_layer=nn.GELU, |
|||
norm_layer=nn.LayerNorm, attn_head_dim=None |
|||
): |
|||
super().__init__() |
|||
|
|||
self.norm1 = norm_layer(dim) |
|||
self.attn = Attention( |
|||
dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, |
|||
attn_drop=attn_drop, proj_drop=drop, attn_head_dim=attn_head_dim |
|||
) |
|||
|
|||
# NOTE: drop path for stochastic depth, we shall see if this is better than dropout here |
|||
self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() |
|||
self.norm2 = norm_layer(dim) |
|||
mlp_hidden_dim = int(dim * mlp_ratio) |
|||
self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) |
|||
|
|||
def forward(self, x): |
|||
x = x + self.drop_path(self.attn(self.norm1(x))) |
|||
x = x + self.drop_path(self.mlp(self.norm2(x))) |
|||
return x |
|||
|
|||
|
|||
class PatchEmbed(nn.Module): |
|||
""" Image to Patch Embedding |
|||
""" |
|||
def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, ratio=1): |
|||
super().__init__() |
|||
img_size = to_2tuple(img_size) |
|||
patch_size = to_2tuple(patch_size) |
|||
num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0]) * (ratio ** 2) |
|||
self.patch_shape = (int(img_size[0] // patch_size[0] * ratio), int(img_size[1] // patch_size[1] * ratio)) |
|||
self.origin_patch_shape = (int(img_size[0] // patch_size[0]), int(img_size[1] // patch_size[1])) |
|||
self.img_size = img_size |
|||
self.patch_size = patch_size |
|||
self.num_patches = num_patches |
|||
|
|||
self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=(patch_size[0] // ratio), padding=4 + 2 * (ratio//2-1)) |
|||
|
|||
def forward(self, x, **kwargs): |
|||
B, C, H, W = x.shape |
|||
x = self.proj(x) |
|||
Hp, Wp = x.shape[2], x.shape[3] |
|||
|
|||
x = x.flatten(2).transpose(1, 2) |
|||
return x, (Hp, Wp) |
|||
|
|||
|
|||
class HybridEmbed(nn.Module): |
|||
""" CNN Feature Map Embedding |
|||
Extract feature map from CNN, flatten, project to embedding dim. |
|||
""" |
|||
def __init__(self, backbone, img_size=224, feature_size=None, in_chans=3, embed_dim=768): |
|||
super().__init__() |
|||
assert isinstance(backbone, nn.Module) |
|||
img_size = to_2tuple(img_size) |
|||
self.img_size = img_size |
|||
self.backbone = backbone |
|||
if feature_size is None: |
|||
with torch.no_grad(): |
|||
training = backbone.training |
|||
if training: |
|||
backbone.eval() |
|||
o = self.backbone(torch.zeros(1, in_chans, img_size[0], img_size[1]))[-1] |
|||
feature_size = o.shape[-2:] |
|||
feature_dim = o.shape[1] |
|||
backbone.train(training) |
|||
else: |
|||
feature_size = to_2tuple(feature_size) |
|||
feature_dim = self.backbone.feature_info.channels()[-1] |
|||
self.num_patches = feature_size[0] * feature_size[1] |
|||
self.proj = nn.Linear(feature_dim, embed_dim) |
|||
|
|||
def forward(self, x): |
|||
x = self.backbone(x)[-1] |
|||
x = x.flatten(2).transpose(1, 2) |
|||
x = self.proj(x) |
|||
return x |
|||
|
|||
class Cross_Attention(nn.Module): |
|||
def __init__(self, dim, num_heads=12, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.): |
|||
super().__init__() |
|||
self.num_heads = num_heads |
|||
head_dim = dim // num_heads |
|||
self.scale = qk_scale or head_dim ** -0.5 |
|||
|
|||
self.self_attn = Attention( |
|||
dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, |
|||
attn_drop=attn_drop, proj_drop=0.) |
|||
|
|||
self.linear_q = nn.Linear(dim, dim, bias=qkv_bias) |
|||
self.linear_k = nn.Linear(dim, dim, bias=qkv_bias) |
|||
self.linear_v = nn.Linear(dim, dim, bias=qkv_bias) |
|||
|
|||
self.attn_drop = nn.Dropout(attn_drop) |
|||
self.proj = nn.Linear(dim, dim) |
|||
self.proj_drop = nn.Dropout(proj_drop) |
|||
|
|||
def forward(self, x_1, x_2, x_3): |
|||
B, N, C = x_1.shape # q |
|||
B, N_1, C = x_2.shape # k, v |
|||
|
|||
q = self.linear_q(x_1).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3) # (B, num_heads, N, C//num_heads) |
|||
k = self.linear_k(x_2).reshape(B, N_1, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3) # (B, num_heads, N_1, C//num_heads) |
|||
v = self.linear_v(x_3).reshape(B, N_1, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3) # (B, num_heads, N_1, C//num_heads) |
|||
|
|||
attn = (q @ k.transpose(-2, -1)) * self.scale # (B, num_heads, N, N_1) |
|||
attn = attn.softmax(dim=-1) |
|||
attn = self.attn_drop(attn) |
|||
|
|||
# import matplotlib.pyplot as plt |
|||
# import seaborn as sns |
|||
|
|||
# attn_map = attn[0][0].cpu().detach().numpy() |
|||
# plt.figure(figsize=(20, 10)) |
|||
# sns.heatmap(attn_map, annot=True, fmt='.2f', cmap='coolwarm') |
|||
|
|||
# plt.title('Cross Attention Map') |
|||
# plt.xlabel('N_1') |
|||
# plt.ylabel('N') |
|||
|
|||
# plt.savefig('/home/fhw/code/ViTPose/test/cross_attn_map.png') |
|||
# plt.close() |
|||
|
|||
x = (attn @ v).transpose(1, 2).reshape(B, N, C) # (B, N, C) |
|||
x = self.proj(x) |
|||
x = self.proj_drop(x) |
|||
return x |
|||
|
|||
class CustomAttentionFFN(nn.Module): |
|||
def __init__(self, dim, num_heads=12, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.): |
|||
super().__init__() |
|||
self.self_attn = Attention( |
|||
dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, |
|||
attn_drop=attn_drop, proj_drop=proj_drop) |
|||
|
|||
self.cross_attn = Cross_Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias, \ |
|||
qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=proj_drop) |
|||
|
|||
self.ffn = nn.Sequential( |
|||
nn.Linear(dim, dim * 4), |
|||
nn.GELU(), |
|||
nn.Linear(dim * 4, dim) |
|||
) |
|||
self.norm1 = nn.LayerNorm(dim) |
|||
self.norm2 = nn.LayerNorm(dim) |
|||
self.norm3 = nn.LayerNorm(dim) |
|||
|
|||
def forward(self, x1, x2): |
|||
x1 = self.norm1(x1 + self.self_attn(x1)) |
|||
|
|||
x1 = self.norm2(x1 + self.cross_attn(x1, x2, x2)) |
|||
|
|||
x1 = self.norm3(x1 + self.ffn(x1)) |
|||
|
|||
return x1 |
|||
|
|||
@BACKBONES.register_module() |
|||
class ViTSam(BaseBackbone): |
|||
|
|||
def __init__(self, |
|||
img_size=224, patch_size=16, in_chans=3, num_classes=80, embed_dim=768, depth=12, |
|||
num_heads=12, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0., |
|||
drop_path_rate=0., hybrid_backbone=None, norm_layer=None, use_checkpoint=False, |
|||
frozen_stages=-1, ratio=1, last_norm=True, |
|||
patch_padding='pad', freeze_attn=False, freeze_ffn=False, samvit_checkpoint=None |
|||
): |
|||
# Protect mutable default arguments |
|||
super(ViTSam, self).__init__() |
|||
norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6) |
|||
self.num_classes = num_classes |
|||
self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models |
|||
self.frozen_stages = frozen_stages |
|||
self.use_checkpoint = use_checkpoint |
|||
self.patch_padding = patch_padding |
|||
self.freeze_attn = freeze_attn |
|||
self.freeze_ffn = freeze_ffn |
|||
self.depth = depth |
|||
|
|||
if hybrid_backbone is not None: |
|||
self.patch_embed = HybridEmbed( |
|||
hybrid_backbone, img_size=img_size, in_chans=in_chans, embed_dim=embed_dim) |
|||
else: |
|||
self.patch_embed = PatchEmbed( |
|||
img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim, ratio=ratio) |
|||
num_patches = self.patch_embed.num_patches |
|||
|
|||
# since the pretraining model has class token |
|||
self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim)) |
|||
|
|||
dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule |
|||
|
|||
self.blocks = nn.ModuleList([ |
|||
Block( |
|||
dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale, |
|||
drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer, |
|||
) |
|||
for i in range(depth)]) |
|||
|
|||
self.last_norm = norm_layer(embed_dim) if last_norm else nn.Identity() |
|||
|
|||
if self.pos_embed is not None: |
|||
trunc_normal_(self.pos_embed, std=.02) |
|||
|
|||
self._freeze_stages() |
|||
|
|||
# ======================== SAM-ViT ======================== |
|||
self.sam_vit = build_vit_sam(model_name='vit_b', checkpoint=samvit_checkpoint) |
|||
self.sam_vit.eval() |
|||
for param in self.sam_vit.parameters(): |
|||
param.requires_grad = False |
|||
|
|||
# self.cross_attn = Cross_Attention(embed_dim, num_heads=num_heads, qkv_bias=qkv_bias, \ |
|||
# qk_scale=qk_scale, attn_drop=attn_drop_rate, proj_drop=drop_rate) |
|||
|
|||
self.custom_attn_ffn = CustomAttentionFFN(embed_dim, num_heads=num_heads, qkv_bias=qkv_bias, \ |
|||
qk_scale=qk_scale, attn_drop=attn_drop_rate, proj_drop=drop_rate) |
|||
|
|||
def _freeze_stages(self): |
|||
"""Freeze parameters.""" |
|||
if self.frozen_stages >= 0: |
|||
self.patch_embed.eval() |
|||
for param in self.patch_embed.parameters(): |
|||
param.requires_grad = False |
|||
|
|||
for i in range(0, self.frozen_stages): |
|||
m = self.blocks[i] |
|||
m.eval() |
|||
for param in m.parameters(): |
|||
param.requires_grad = False |
|||
|
|||
if self.freeze_attn: |
|||
for i in range(0, self.depth): |
|||
m = self.blocks[i] |
|||
m.attn.eval() |
|||
m.norm1.eval() |
|||
for param in m.attn.parameters(): |
|||
param.requires_grad = False |
|||
for param in m.norm1.parameters(): |
|||
param.requires_grad = False |
|||
|
|||
if self.freeze_ffn: |
|||
self.pos_embed.requires_grad = False |
|||
self.patch_embed.eval() |
|||
for param in self.patch_embed.parameters(): |
|||
param.requires_grad = False |
|||
for i in range(0, self.depth): |
|||
m = self.blocks[i] |
|||
m.mlp.eval() |
|||
m.norm2.eval() |
|||
for param in m.mlp.parameters(): |
|||
param.requires_grad = False |
|||
for param in m.norm2.parameters(): |
|||
param.requires_grad = False |
|||
|
|||
def init_weights(self, pretrained=None): |
|||
"""Initialize the weights in backbone. |
|||
Args: |
|||
pretrained (str, optional): Path to pre-trained weights. |
|||
Defaults to None. |
|||
""" |
|||
super().init_weights(pretrained, patch_padding=self.patch_padding) |
|||
|
|||
if pretrained is None: |
|||
def _init_weights(m): |
|||
if isinstance(m, nn.Linear): |
|||
trunc_normal_(m.weight, std=.02) |
|||
if isinstance(m, nn.Linear) and m.bias is not None: |
|||
nn.init.constant_(m.bias, 0) |
|||
elif isinstance(m, nn.LayerNorm): |
|||
nn.init.constant_(m.bias, 0) |
|||
nn.init.constant_(m.weight, 1.0) |
|||
|
|||
self.apply(_init_weights) |
|||
|
|||
def get_num_layers(self): |
|||
return len(self.blocks) |
|||
|
|||
@torch.jit.ignore |
|||
def no_weight_decay(self): |
|||
return {'pos_embed', 'cls_token'} |
|||
|
|||
def forward_features(self, x): |
|||
B, C, H, W = x.shape |
|||
x, (Hp, Wp) = self.patch_embed(x) |
|||
|
|||
if self.pos_embed is not None: |
|||
# fit for multiple GPU training |
|||
# since the first element for pos embed (sin-cos manner) is zero, it will cause no difference |
|||
x = x + self.pos_embed[:, 1:] + self.pos_embed[:, :1] |
|||
|
|||
for blk in self.blocks: |
|||
if self.use_checkpoint: |
|||
x = checkpoint.checkpoint(blk, x) |
|||
else: |
|||
x = blk(x) |
|||
|
|||
x = self.last_norm(x) |
|||
|
|||
return x, Hp, Wp |
|||
|
|||
def forward(self, x1, x2): |
|||
import time |
|||
B, _, _, _ = x1.shape |
|||
x1, Hp, Wp = self.forward_features(x1) # B, N_vitpose, C |
|||
|
|||
with torch.no_grad(): |
|||
# start_time = time.time() |
|||
# self.sam_vit.eval() |
|||
x2 = self.sam_vit(x2) # B, N_sam, C |
|||
|
|||
# end_time = time.time() |
|||
# print('SAM-ViT forward time: {:.4f}秒'.format(end_time - start_time)) |
|||
|
|||
# x1 = x1 + self.cross_attn(x1, x2, x2) |
|||
|
|||
x1 = self.custom_attn_ffn(x1, x2) |
|||
|
|||
xp = x1.permute(0, 2, 1).reshape(B, -1, Hp, Wp).contiguous() # B, C, Hp, Wp |
|||
return xp |
|||
|
|||
def train(self, mode=True): |
|||
"""Convert the model into training mode.""" |
|||
super().train(mode) |
|||
self._freeze_stages() |
|||
|
|||
|
|||
if __name__ == '__main__': |
|||
from thop import profile |
|||
from mmcv.runner import load_checkpoint |
|||
|
|||
# Prepare an image as the input |
|||
bs, c, h, w = 2, 3, 1024, 1024 |
|||
x1 = torch.randn(bs, c, 256, 192) |
|||
x2 = torch.randn(bs, c, h, w) |
|||
patch_size = 16 |
|||
device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu') |
|||
|
|||
# Build model |
|||
model = ViTSam(img_size=(256, 192), patch_size=16, embed_dim=768, depth=12, num_heads=12, ratio=1, |
|||
use_checkpoint=False, mlp_ratio=4, qkv_bias=True, drop_path_rate=0.3, |
|||
samvit_checkpoint='/home/fhw/code/ViTPose/checkpoints/sam/sam_vit_b_01ec64.pth') |
|||
|
|||
|
|||
if torch.cuda.is_available(): |
|||
x1 = x1.to(device) |
|||
x2 = x2.to(device) |
|||
model = model.to(device) |
|||
|
|||
with torch.no_grad(): |
|||
model.eval() |
|||
# Inference |
|||
outputs = model(x1, x2) |
|||
print(outputs.shape) |
|||
|
|||
# Compute FLOPs & Params |
|||
print('==============================') |
|||
model.eval() |
|||
flops, params = profile(model, inputs=(x1, x2), verbose=False) |
|||
print('GFLOPs : {:.2f}'.format(flops / 1e9 * 2)) |
|||
print('Params : {:.2f} M'.format(params / 1e6)) |
|||
@ -0,0 +1,322 @@ |
|||
# Copyright (c) OpenMMLab. All rights reserved. |
|||
import warnings |
|||
import logging |
|||
|
|||
import mmcv |
|||
import numpy as np |
|||
from mmcv.image import imwrite |
|||
from mmcv.utils.misc import deprecated_api_warning |
|||
from mmcv.visualization.image import imshow |
|||
from mmcv_custom.checkpoint import load_checkpoint |
|||
|
|||
from mmpose.core import imshow_bboxes, imshow_keypoints |
|||
from .. import builder |
|||
from ..builder import POSENETS |
|||
from .base import BasePose |
|||
|
|||
try: |
|||
from mmcv.runner import auto_fp16 |
|||
except ImportError: |
|||
warnings.warn('auto_fp16 from mmpose will be deprecated from v0.15.0' |
|||
'Please install mmcv>=1.1.4') |
|||
from mmpose.core import auto_fp16 |
|||
|
|||
|
|||
@POSENETS.register_module() |
|||
class TopDownSelf(BasePose): |
|||
"""Top-down pose detectors. |
|||
|
|||
Args: |
|||
backbone (dict): Backbone modules to extract feature. |
|||
keypoint_head (dict): Keypoint head to process feature. |
|||
train_cfg (dict): Config for training. Default: None. |
|||
test_cfg (dict): Config for testing. Default: None. |
|||
pretrained (str): Path to the pretrained models. |
|||
loss_pose (None): Deprecated arguments. Please use |
|||
`loss_keypoint` for heads instead. |
|||
""" |
|||
|
|||
def __init__(self, |
|||
backbone, |
|||
neck=None, |
|||
keypoint_head=None, |
|||
train_cfg=None, |
|||
test_cfg=None, |
|||
pretrained=None, |
|||
loss_pose=None): |
|||
super().__init__() |
|||
self.fp16_enabled = False |
|||
|
|||
self.backbone = builder.build_backbone(backbone) |
|||
|
|||
self.train_cfg = train_cfg |
|||
self.test_cfg = test_cfg |
|||
|
|||
if neck is not None: |
|||
self.neck = builder.build_neck(neck) |
|||
|
|||
if keypoint_head is not None: |
|||
keypoint_head['train_cfg'] = train_cfg |
|||
keypoint_head['test_cfg'] = test_cfg |
|||
|
|||
if 'loss_keypoint' not in keypoint_head and loss_pose is not None: |
|||
warnings.warn( |
|||
'`loss_pose` for TopDown is deprecated, ' |
|||
'use `loss_keypoint` for heads instead. See ' |
|||
'https://github.com/open-mmlab/mmpose/pull/382' |
|||
' for more information.', DeprecationWarning) |
|||
keypoint_head['loss_keypoint'] = loss_pose |
|||
|
|||
self.keypoint_head = builder.build_head(keypoint_head) |
|||
|
|||
self.init_weights(pretrained=pretrained) |
|||
|
|||
@property |
|||
def with_neck(self): |
|||
"""Check if has neck.""" |
|||
return hasattr(self, 'neck') |
|||
|
|||
@property |
|||
def with_keypoint(self): |
|||
"""Check if has keypoint_head.""" |
|||
return hasattr(self, 'keypoint_head') |
|||
|
|||
def init_weights(self, pretrained=None): |
|||
"""Weight initialization for model.""" |
|||
self.backbone.init_weights(pretrained) |
|||
if self.with_neck: |
|||
self.neck.init_weights() |
|||
if self.with_keypoint: |
|||
self.keypoint_head.init_weights() |
|||
|
|||
@auto_fp16(apply_to=('img', 'sam_img', )) |
|||
def forward(self, |
|||
img, |
|||
sam_img, # 针对sam_encoder的输入 |
|||
target=None, |
|||
target_weight=None, |
|||
img_metas=None, |
|||
return_loss=True, |
|||
return_heatmap=False, |
|||
**kwargs): |
|||
"""Calls either forward_train or forward_test depending on whether |
|||
return_loss=True. Note this setting will change the expected inputs. |
|||
When `return_loss=True`, img and img_meta are single-nested (i.e. |
|||
Tensor and List[dict]), and when `resturn_loss=False`, img and img_meta |
|||
should be double nested (i.e. List[Tensor], List[List[dict]]), with |
|||
the outer list indicating test time augmentations. |
|||
|
|||
Note: |
|||
- batch_size: N |
|||
- num_keypoints: K |
|||
- num_img_channel: C (Default: 3) |
|||
- img height: imgH |
|||
- img width: imgW |
|||
- heatmaps height: H |
|||
- heatmaps weight: W |
|||
|
|||
Args: |
|||
img (torch.Tensor[NxCximgHximgW]): Input images. |
|||
target (torch.Tensor[NxKxHxW]): Target heatmaps. |
|||
target_weight (torch.Tensor[NxKx1]): Weights across |
|||
different joint types. |
|||
img_metas (list(dict)): Information about data augmentation |
|||
By default this includes: |
|||
|
|||
- "image_file: path to the image file |
|||
- "center": center of the bbox |
|||
- "scale": scale of the bbox |
|||
- "rotation": rotation of the bbox |
|||
- "bbox_score": score of bbox |
|||
return_loss (bool): Option to `return loss`. `return loss=True` |
|||
for training, `return loss=False` for validation & test. |
|||
return_heatmap (bool) : Option to return heatmap. |
|||
|
|||
Returns: |
|||
dict|tuple: if `return loss` is true, then return losses. \ |
|||
Otherwise, return predicted poses, boxes, image paths \ |
|||
and heatmaps. |
|||
""" |
|||
if return_loss: |
|||
# 可视化 img, sam_img cv可视化/PIL Image |
|||
# print(sam_img[0].shape) |
|||
# imshow(sam_img[0].cpu().numpy().transpose(1, 2, 0), wait_time=5000) |
|||
# 修改 |
|||
return self.forward_train(img, sam_img, target, target_weight, img_metas, |
|||
**kwargs) |
|||
# 修改 |
|||
return self.forward_test( |
|||
img, sam_img, img_metas, return_heatmap=return_heatmap, **kwargs) |
|||
|
|||
# 修改 |
|||
def forward_train(self, img, sam_img, target, target_weight, img_metas, **kwargs): |
|||
"""Defines the computation performed at every call when training.""" |
|||
# 修改 |
|||
output = self.backbone(img, sam_img) # B, C, Hp, Wp |
|||
if self.with_neck: |
|||
output = self.neck(output) |
|||
if self.with_keypoint: |
|||
output = self.keypoint_head(output) |
|||
|
|||
# if return loss |
|||
losses = dict() |
|||
if self.with_keypoint: |
|||
keypoint_losses = self.keypoint_head.get_loss( |
|||
output, target, target_weight) |
|||
losses.update(keypoint_losses) |
|||
keypoint_accuracy = self.keypoint_head.get_accuracy( |
|||
output, target, target_weight) |
|||
losses.update(keypoint_accuracy) |
|||
|
|||
return losses |
|||
|
|||
# 修改 |
|||
def forward_test(self, img, sam_img, img_metas, return_heatmap=False, **kwargs): |
|||
"""Defines the computation performed at every call when testing.""" |
|||
assert img.size(0) == len(img_metas) |
|||
batch_size, _, img_height, img_width = img.shape |
|||
if batch_size > 1: |
|||
assert 'bbox_id' in img_metas[0] |
|||
|
|||
result = {} |
|||
|
|||
# 修改 |
|||
features = self.backbone(img, sam_img) |
|||
if self.with_neck: |
|||
features = self.neck(features) |
|||
if self.with_keypoint: |
|||
output_heatmap = self.keypoint_head.inference_model( |
|||
features, flip_pairs=None) |
|||
|
|||
if self.test_cfg.get('flip_test', True): |
|||
img_flipped = img.flip(3) |
|||
# 修改 |
|||
sam_img_flipped = sam_img.flip(3) |
|||
features_flipped = self.backbone(img_flipped, sam_img_flipped) |
|||
if self.with_neck: |
|||
features_flipped = self.neck(features_flipped) |
|||
if self.with_keypoint: |
|||
output_flipped_heatmap = self.keypoint_head.inference_model( |
|||
features_flipped, img_metas[0]['flip_pairs']) |
|||
output_heatmap = (output_heatmap + |
|||
output_flipped_heatmap) * 0.5 |
|||
|
|||
if self.with_keypoint: |
|||
keypoint_result = self.keypoint_head.decode( |
|||
img_metas, output_heatmap, img_size=[img_width, img_height]) |
|||
result.update(keypoint_result) |
|||
|
|||
if not return_heatmap: |
|||
output_heatmap = None |
|||
|
|||
result['output_heatmap'] = output_heatmap |
|||
|
|||
return result |
|||
|
|||
# 修改 |
|||
def forward_dummy(self, img, sam_img): |
|||
"""Used for computing network FLOPs. |
|||
|
|||
See ``tools/get_flops.py``. |
|||
|
|||
Args: |
|||
img (torch.Tensor): Input image. |
|||
|
|||
Returns: |
|||
Tensor: Output heatmaps. |
|||
""" |
|||
output = self.backbone(img, sam_img) |
|||
if self.with_neck: |
|||
output = self.neck(output) |
|||
if self.with_keypoint: |
|||
output = self.keypoint_head(output) |
|||
return output |
|||
|
|||
@deprecated_api_warning({'pose_limb_color': 'pose_link_color'}, |
|||
cls_name='TopDown') |
|||
def show_result(self, |
|||
img, |
|||
result, |
|||
skeleton=None, |
|||
kpt_score_thr=0.3, |
|||
bbox_color='green', |
|||
pose_kpt_color=None, |
|||
pose_link_color=None, |
|||
text_color='white', |
|||
radius=4, |
|||
thickness=1, |
|||
font_scale=0.5, |
|||
bbox_thickness=1, |
|||
win_name='', |
|||
show=False, |
|||
show_keypoint_weight=False, |
|||
wait_time=0, |
|||
out_file=None): |
|||
"""Draw `result` over `img`. |
|||
|
|||
Args: |
|||
img (str or Tensor): The image to be displayed. |
|||
result (list[dict]): The results to draw over `img` |
|||
(bbox_result, pose_result). |
|||
skeleton (list[list]): The connection of keypoints. |
|||
skeleton is 0-based indexing. |
|||
kpt_score_thr (float, optional): Minimum score of keypoints |
|||
to be shown. Default: 0.3. |
|||
bbox_color (str or tuple or :obj:`Color`): Color of bbox lines. |
|||
pose_kpt_color (np.array[Nx3]`): Color of N keypoints. |
|||
If None, do not draw keypoints. |
|||
pose_link_color (np.array[Mx3]): Color of M links. |
|||
If None, do not draw links. |
|||
text_color (str or tuple or :obj:`Color`): Color of texts. |
|||
radius (int): Radius of circles. |
|||
thickness (int): Thickness of lines. |
|||
font_scale (float): Font scales of texts. |
|||
win_name (str): The window name. |
|||
show (bool): Whether to show the image. Default: False. |
|||
show_keypoint_weight (bool): Whether to change the transparency |
|||
using the predicted confidence scores of keypoints. |
|||
wait_time (int): Value of waitKey param. |
|||
Default: 0. |
|||
out_file (str or None): The filename to write the image. |
|||
Default: None. |
|||
|
|||
Returns: |
|||
Tensor: Visualized img, only if not `show` or `out_file`. |
|||
""" |
|||
img = mmcv.imread(img) |
|||
img = img.copy() |
|||
|
|||
bbox_result = [] |
|||
bbox_labels = [] |
|||
pose_result = [] |
|||
for res in result: |
|||
if 'bbox' in res: |
|||
bbox_result.append(res['bbox']) |
|||
bbox_labels.append(res.get('label', None)) |
|||
pose_result.append(res['keypoints']) |
|||
|
|||
if bbox_result: |
|||
bboxes = np.vstack(bbox_result) |
|||
# draw bounding boxes |
|||
imshow_bboxes( |
|||
img, |
|||
bboxes, |
|||
labels=bbox_labels, |
|||
colors=bbox_color, |
|||
text_color=text_color, |
|||
thickness=bbox_thickness, |
|||
font_scale=font_scale, |
|||
show=False) |
|||
|
|||
if pose_result: |
|||
imshow_keypoints(img, pose_result, skeleton, kpt_score_thr, |
|||
pose_kpt_color, pose_link_color, radius, |
|||
thickness) |
|||
|
|||
if show: |
|||
imshow(img, win_name, wait_time) |
|||
|
|||
if out_file is not None: |
|||
imwrite(img, out_file) |
|||
|
|||
return img |
|||
@ -0,0 +1,17 @@ |
|||
import torch |
|||
import numpy as np |
|||
|
|||
model_1 = torch.load('/home/fhw/code/ViTPose/checkpoints/sam/sam_vit_b_01ec64.pth') |
|||
model_2 = torch.load('/home/fhw/code/ViTPose/work_dirs/ViTSam_base_coco_256x192/best_AP_epoch_1.pth') |
|||
|
|||
param_1 = model_1['image_encoder.pos_embed'].numpy() |
|||
param_2 = model_2['state_dict']['backbone.sam_vit.pos_embed'].numpy() |
|||
|
|||
# for name, param in model_2.items(): |
|||
# print(name) |
|||
|
|||
# print(model_2['state_dict']['backbone.sam_vit.pos_embed']) |
|||
|
|||
is_equal = np.array_equal(param_1, param_2) |
|||
|
|||
print(is_equal) |
|||
@ -0,0 +1,197 @@ |
|||
# Copyright (c) OpenMMLab. All rights reserved. |
|||
import argparse |
|||
import copy |
|||
import os |
|||
import os.path as osp |
|||
import time |
|||
import warnings |
|||
|
|||
import mmcv |
|||
import torch |
|||
from mmcv import Config, DictAction |
|||
from mmcv.runner import get_dist_info, init_dist, set_random_seed, load_checkpoint |
|||
from mmcv.utils import get_git_hash |
|||
|
|||
from mmpose import __version__ |
|||
from mmpose.apis import init_random_seed, train_model |
|||
from mmpose.datasets import build_dataset |
|||
from mmpose.models import build_posenet |
|||
from mmpose.utils import collect_env, get_root_logger, setup_multi_processes |
|||
import mmcv_custom |
|||
|
|||
def parse_args(): |
|||
parser = argparse.ArgumentParser(description='Train a pose model') |
|||
parser.add_argument('config', help='train config file path') |
|||
parser.add_argument('-c', '--checkpoint', help='checkpoint file', default='/root/autodl-tmp/code/ViTPose/checkpoints/vitpose/vitpose-b.pth') |
|||
parser.add_argument('--work-dir', help='the dir to save logs and models') |
|||
parser.add_argument( |
|||
'--resume-from', help='the checkpoint file to resume from') |
|||
parser.add_argument( |
|||
'--no-validate', |
|||
action='store_true', |
|||
help='whether not to evaluate the checkpoint during training') |
|||
group_gpus = parser.add_mutually_exclusive_group() |
|||
group_gpus.add_argument( |
|||
'--gpus', |
|||
type=int, |
|||
help='(Deprecated, please use --gpu-id) number of gpus to use ' |
|||
'(only applicable to non-distributed training)') |
|||
group_gpus.add_argument( |
|||
'--gpu-ids', |
|||
type=int, |
|||
nargs='+', |
|||
help='(Deprecated, please use --gpu-id) ids of gpus to use ' |
|||
'(only applicable to non-distributed training)') |
|||
group_gpus.add_argument( |
|||
'--gpu-id', |
|||
type=int, |
|||
default=0, |
|||
help='id of gpu to use ' |
|||
'(only applicable to non-distributed training)') |
|||
parser.add_argument('--seed', type=int, default=None, help='random seed') |
|||
parser.add_argument( |
|||
'--deterministic', |
|||
action='store_true', |
|||
help='whether to set deterministic options for CUDNN backend.') |
|||
parser.add_argument( |
|||
'--cfg-options', |
|||
nargs='+', |
|||
action=DictAction, |
|||
default={}, |
|||
help='override some settings in the used config, the key-value pair ' |
|||
'in xxx=yyy format will be merged into config file. For example, ' |
|||
"'--cfg-options model.backbone.depth=18 model.backbone.with_cp=True'") |
|||
parser.add_argument( |
|||
'--launcher', |
|||
choices=['none', 'pytorch', 'slurm', 'mpi'], |
|||
default='none', |
|||
help='job launcher') |
|||
parser.add_argument('--local_rank', type=int, default=0) |
|||
parser.add_argument( |
|||
'--autoscale-lr', |
|||
action='store_true', |
|||
help='automatically scale lr with the number of gpus') |
|||
args = parser.parse_args() |
|||
if 'LOCAL_RANK' not in os.environ: |
|||
os.environ['LOCAL_RANK'] = str(args.local_rank) |
|||
|
|||
return args |
|||
|
|||
|
|||
def main(): |
|||
args = parse_args() |
|||
|
|||
cfg = Config.fromfile(args.config) |
|||
|
|||
if args.cfg_options is not None: |
|||
cfg.merge_from_dict(args.cfg_options) |
|||
|
|||
# set multi-process settings |
|||
setup_multi_processes(cfg) |
|||
|
|||
# set cudnn_benchmark |
|||
if cfg.get('cudnn_benchmark', False): |
|||
torch.backends.cudnn.benchmark = True |
|||
|
|||
# work_dir is determined in this priority: CLI > segment in file > filename |
|||
if args.work_dir is not None: |
|||
# update configs according to CLI args if args.work_dir is not None |
|||
cfg.work_dir = args.work_dir |
|||
elif cfg.get('work_dir', None) is None: |
|||
# use config filename as default work_dir if cfg.work_dir is None |
|||
cfg.work_dir = osp.join('./work_dirs', |
|||
osp.splitext(osp.basename(args.config))[0]) |
|||
if args.resume_from is not None: |
|||
cfg.resume_from = args.resume_from |
|||
if args.gpus is not None: |
|||
cfg.gpu_ids = range(1) |
|||
warnings.warn('`--gpus` is deprecated because we only support ' |
|||
'single GPU mode in non-distributed training. ' |
|||
'Use `gpus=1` now.') |
|||
if args.gpu_ids is not None: |
|||
cfg.gpu_ids = args.gpu_ids[0:1] |
|||
warnings.warn('`--gpu-ids` is deprecated, please use `--gpu-id`. ' |
|||
'Because we only support single GPU mode in ' |
|||
'non-distributed training. Use the first GPU ' |
|||
'in `gpu_ids` now.') |
|||
if args.gpus is None and args.gpu_ids is None: |
|||
cfg.gpu_ids = [args.gpu_id] |
|||
|
|||
if args.autoscale_lr: |
|||
# apply the linear scaling rule (https://arxiv.org/abs/1706.02677) |
|||
cfg.optimizer['lr'] = cfg.optimizer['lr'] * len(cfg.gpu_ids) / 8 |
|||
|
|||
# init distributed env first, since logger depends on the dist info. |
|||
if args.launcher == 'none': |
|||
distributed = False |
|||
if len(cfg.gpu_ids) > 1: |
|||
warnings.warn( |
|||
f'We treat {cfg.gpu_ids} as gpu-ids, and reset to ' |
|||
f'{cfg.gpu_ids[0:1]} as gpu-ids to avoid potential error in ' |
|||
'non-distribute training time.') |
|||
cfg.gpu_ids = cfg.gpu_ids[0:1] |
|||
else: |
|||
distributed = True |
|||
init_dist(args.launcher, **cfg.dist_params) |
|||
# re-set gpu_ids with distributed training mode |
|||
_, world_size = get_dist_info() |
|||
cfg.gpu_ids = range(world_size) |
|||
|
|||
# create work_dir |
|||
mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir)) |
|||
# init the logger before other steps |
|||
timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime()) |
|||
log_file = osp.join(cfg.work_dir, f'{timestamp}.log') |
|||
logger = get_root_logger(log_file=log_file, log_level=cfg.log_level) |
|||
|
|||
# init the meta dict to record some important information such as |
|||
# environment info and seed, which will be logged |
|||
meta = dict() |
|||
# log env info |
|||
env_info_dict = collect_env() |
|||
env_info = '\n'.join([(f'{k}: {v}') for k, v in env_info_dict.items()]) |
|||
dash_line = '-' * 60 + '\n' |
|||
logger.info('Environment info:\n' + dash_line + env_info + '\n' + |
|||
dash_line) |
|||
meta['env_info'] = env_info |
|||
|
|||
# log some basic info |
|||
logger.info(f'Distributed training: {distributed}') |
|||
logger.info(f'Config:\n{cfg.pretty_text}') |
|||
|
|||
# set random seeds |
|||
seed = init_random_seed(args.seed) |
|||
logger.info(f'Set random seed to {seed}, ' |
|||
f'deterministic: {args.deterministic}') |
|||
set_random_seed(seed, deterministic=args.deterministic) |
|||
cfg.seed = seed |
|||
meta['seed'] = seed |
|||
|
|||
model = build_posenet(cfg.model) |
|||
load_checkpoint(model, args.checkpoint, map_location='cpu') |
|||
datasets = [build_dataset(cfg.data.train)] |
|||
|
|||
if len(cfg.workflow) == 2: |
|||
val_dataset = copy.deepcopy(cfg.data.val) |
|||
val_dataset.pipeline = cfg.data.train.pipeline |
|||
datasets.append(build_dataset(val_dataset)) |
|||
|
|||
if cfg.checkpoint_config is not None: |
|||
# save mmpose version, config file content |
|||
# checkpoints as meta data |
|||
cfg.checkpoint_config.meta = dict( |
|||
mmpose_version=__version__ + get_git_hash(digits=7), |
|||
config=cfg.pretty_text, |
|||
) |
|||
train_model( |
|||
model, |
|||
datasets, |
|||
cfg, |
|||
distributed=distributed, |
|||
validate=(not args.no_validate), |
|||
timestamp=timestamp, |
|||
meta=meta) |
|||
|
|||
|
|||
if __name__ == '__main__': |
|||
main() |
|||
File diff suppressed because it is too large
@ -0,0 +1 @@ |
|||
安装环境时的问题：由于setuptools版本过高，导致算法使用的安装方式已经被弃用，建议选择重新安装小于60的版本 |
|||
Loading…
Reference in new issue