|
|
|
|
|
|
|
|
|
|
|
|
|
|
import argparse
|
|
|
|
|
|
import numpy as np
|
|
|
import onnxruntime
|
|
|
import torch
|
|
|
import torch.nn as nn
|
|
|
from mmengine import Config
|
|
|
from mmengine.registry import init_default_scope
|
|
|
from mmengine.runner import load_checkpoint
|
|
|
from mmengine.structures import LabelData
|
|
|
|
|
|
from mmaction.registry import MODELS
|
|
|
from mmaction.structures import ActionDataSample
|
|
|
|
|
|
|
|
|
def parse_args():
|
|
|
parser = argparse.ArgumentParser(description='Get model flops and params')
|
|
|
parser.add_argument('config', help='config file path')
|
|
|
parser.add_argument('checkpoint', help='checkpoint file')
|
|
|
parser.add_argument(
|
|
|
'--num_frames', type=int, default=48, help='number of input frames.')
|
|
|
parser.add_argument(
|
|
|
'--image_size', type=int, default=64, help='size of the frame')
|
|
|
parser.add_argument(
|
|
|
'--num_joints',
|
|
|
type=int,
|
|
|
default=0,
|
|
|
help='number of joints. If not given, will use default settings from'
|
|
|
'the config file')
|
|
|
parser.add_argument(
|
|
|
'--device', type=str, default='cpu', help='CPU/CUDA device option')
|
|
|
parser.add_argument(
|
|
|
'--output_file',
|
|
|
type=str,
|
|
|
default='posec3d.onnx',
|
|
|
help='file name of the output onnx file')
|
|
|
args = parser.parse_args()
|
|
|
return args
|
|
|
|
|
|
|
|
|
class AvgPool3d(nn.Module):
|
|
|
|
|
|
def forward(self, x):
|
|
|
return x.mean(dim=(-1, -2, -3), keepdims=True)
|
|
|
|
|
|
|
|
|
class MaxPool3d(nn.Module):
|
|
|
|
|
|
def forward(self, x):
|
|
|
x = x.max(dim=-1, keepdim=True)[0]
|
|
|
x = x.max(dim=-2, keepdim=True)[0]
|
|
|
x = x.max(dim=-3, keepdim=True)[0]
|
|
|
return x
|
|
|
|
|
|
|
|
|
class GCNNet(nn.Module):
|
|
|
|
|
|
def __init__(self, base_model):
|
|
|
super(GCNNet, self).__init__()
|
|
|
self.backbone = base_model.backbone
|
|
|
self.head = base_model.cls_head
|
|
|
|
|
|
if hasattr(self.head, 'pool'):
|
|
|
pool = self.head.pool
|
|
|
if isinstance(pool, nn.AdaptiveAvgPool3d):
|
|
|
assert pool.output_size == 1
|
|
|
self.head.pool = AvgPool3d()
|
|
|
elif isinstance(pool, nn.AdaptiveMaxPool3d):
|
|
|
assert pool.output_size == 1
|
|
|
self.head.pool = MaxPool3d()
|
|
|
|
|
|
def forward(self, input_tensor):
|
|
|
feat = self.backbone(input_tensor)
|
|
|
cls_score = self.head(feat)
|
|
|
return cls_score
|
|
|
|
|
|
|
|
|
def softmax(x):
|
|
|
x = np.exp(x - x.max())
|
|
|
return x / x.sum()
|
|
|
|
|
|
|
|
|
def main():
|
|
|
args = parse_args()
|
|
|
config = Config.fromfile(args.config)
|
|
|
|
|
|
if config.model.type != 'RecognizerGCN':
|
|
|
print('This script serves the sole purpose of converting PoseC3D '
|
|
|
'skeleton models in MMAction2 to ONNX files. Please note that '
|
|
|
'attempting to convert other models using this script may not '
|
|
|
'yield successful results.\n\n')
|
|
|
|
|
|
init_default_scope(config.get('default_scope', 'mmaction'))
|
|
|
|
|
|
base_model = MODELS.build(config.model)
|
|
|
load_checkpoint(base_model, args.checkpoint, map_location='cpu')
|
|
|
base_model.to(args.device)
|
|
|
|
|
|
num_joints = args.num_joints
|
|
|
image_size = args.image_size
|
|
|
num_frames = args.num_frames
|
|
|
if num_joints == 0:
|
|
|
num_joints = config.model.backbone.in_channels
|
|
|
|
|
|
input_tensor = torch.randn(1, num_joints, num_frames, image_size,
|
|
|
image_size)
|
|
|
input_tensor = input_tensor.clamp(-3, 3).to(args.device)
|
|
|
|
|
|
base_model.eval()
|
|
|
|
|
|
data_sample = ActionDataSample()
|
|
|
data_sample.pred_scores = LabelData()
|
|
|
data_sample.pred_labels = LabelData()
|
|
|
base_output = base_model(
|
|
|
input_tensor.unsqueeze(0), data_samples=[data_sample],
|
|
|
mode='predict')[0]
|
|
|
base_output = base_output.pred_score.detach().cpu().numpy()
|
|
|
|
|
|
model = GCNNet(base_model).to(args.device)
|
|
|
model.eval()
|
|
|
|
|
|
torch.onnx.export(
|
|
|
model, (input_tensor),
|
|
|
args.output_file,
|
|
|
input_names=['input_tensor'],
|
|
|
output_names=['cls_score'],
|
|
|
export_params=True,
|
|
|
do_constant_folding=True,
|
|
|
verbose=False,
|
|
|
opset_version=11,
|
|
|
dynamic_axes={
|
|
|
'input_tensor': {
|
|
|
0: 'batch_size',
|
|
|
2: 'num_frames'
|
|
|
},
|
|
|
'cls_score': {
|
|
|
0: 'batch_size'
|
|
|
}
|
|
|
})
|
|
|
|
|
|
print(f'Successfully export the onnx file to {args.output_file}')
|
|
|
|
|
|
|
|
|
session = onnxruntime.InferenceSession(args.output_file)
|
|
|
input_feed = {'input_tensor': input_tensor.cpu().data.numpy()}
|
|
|
outputs = session.run(['cls_score'], input_feed=input_feed)
|
|
|
output = softmax(outputs[0][0])
|
|
|
|
|
|
diff = abs(base_output - output).max()
|
|
|
if diff < 1e-5:
|
|
|
print('The output difference is smaller than 1e-5.')
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
main()
|
|
|
|
