NVIDIA-AI-IOT / trt_pose

Real-time pose estimation accelerated with NVIDIA TensorRT
MIT License
982 stars 293 forks source link

I changed backbone to mobilenet, but the speed is very slow. #140

Closed sinuku closed 3 years ago

sinuku commented 3 years ago

Hi! Thanks for you project.

I changed backbone to mobilenet, but the speed is very slow.

I created my_mobilenetv2.py in the models folder and modified it as follows.

class MobileNetV2Backbone(torch.nn.Module):
    def __init__(self, mobilenetv2):
        super(MobileNetV2Backbone, self).__init__()
        self.mobilenetv2 = mobilenetv2

    def forward(self, x):
        x = self.mobilenetv2.features(x)

        return x

def _mobilenetv2_pose_att(cmap_channels, paf_channels, upsample_channels, mobilenetv2, feature_channels, num_upsample, num_flat):
    model = torch.nn.Sequential(
        MobileNetV2Backbone(mobilenetv2),
        CmapPafHeadAttention(feature_channels, cmap_channels, paf_channels, upsample_channels, num_upsample=num_upsample, num_flat=num_flat)
    )
    return model

def my_mobilenetv2_baseline_att(cmap_channels, paf_channels, upsample_channels=256, pretrained=True, num_upsample=3, num_flat=0):
    mobilenetv2 = mobilenet_v2(pretrained=pretrained)
    return _mobilenetv2_pose_att(cmap_channels, paf_channels, upsample_channels, mobilenetv2, 1280, num_upsample, num_flat)

The config file follows the format of resnet18_baseline_att_224x224_A.json What's the problem?

sinuku commented 3 years ago

resolved the issue