Shortcuts

Source code for mmpose.models.necks.gap_neck

# Copyright (c) OpenMMLab. All rights reserved.
import torch
import torch.nn as nn

from mmpose.registry import MODELS


[docs]@MODELS.register_module() class GlobalAveragePooling(nn.Module): """Global Average Pooling neck. Note that we use `view` to remove extra channel after pooling. We do not use `squeeze` as it will also remove the batch dimension when the tensor has a batch dimension of size 1, which can lead to unexpected errors. """ def __init__(self): super().__init__() self.gap = nn.AdaptiveAvgPool2d((1, 1)) def init_weights(self): pass
[docs] def forward(self, inputs): """Forward function.""" if isinstance(inputs, tuple): outs = tuple([self.gap(x) for x in inputs]) outs = tuple( [out.view(x.size(0), -1) for out, x in zip(outs, inputs)]) elif isinstance(inputs, list): outs = [self.gap(x) for x in inputs] outs = [out.view(x.size(0), -1) for out, x in zip(outs, inputs)] elif isinstance(inputs, torch.Tensor): outs = self.gap(inputs) outs = outs.view(inputs.size(0), -1) else: raise TypeError('neck inputs should be tuple or torch.tensor') return outs
Read the Docs v: latest
Versions
latest
0.x
dev-1.x
Downloads
epub
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.