Skip to content

Commit

Permalink
Merge branch 'ljt/move_avgconsensus' into 'master'
Browse files Browse the repository at this point in the history
Move AvgConsensus to head/base.py

See merge request EIG-Research/mmaction-lite!86
  • Loading branch information
chenkai committed Mar 24, 2020
2 parents 53d90e9 + ed80917 commit 5e84344
Show file tree
Hide file tree
Showing 2 changed files with 17 additions and 17 deletions.
16 changes: 16 additions & 0 deletions mmaction/models/heads/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,22 @@
import torch.nn.functional as F


class AvgConsensus(nn.Module):
"""Average consensus module.
Attributes:
dim (int): Decide which dim consensus function to apply.
Default: 1.
"""

def __init__(self, dim=1):
super(AvgConsensus, self).__init__()
self.dim = dim

def forward(self, input):
return input.mean(dim=self.dim, keepdim=True)


class BaseHead(nn.Module, metaclass=ABCMeta):
"""Base class for head.
Expand Down
18 changes: 1 addition & 17 deletions mmaction/models/heads/tsn_head.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,23 +2,7 @@
from mmcv.cnn.weight_init import normal_init

from ..registry import HEADS
from .base import BaseHead


class AvgConsensus(nn.Module):
"""Average consensus module.
Attributes:
dim (int): Decide which dim consensus function to apply.
Default: 1.
"""

def __init__(self, dim=1):
super(AvgConsensus, self).__init__()
self.dim = dim

def forward(self, input):
return input.mean(dim=self.dim, keepdim=True)
from .base import AvgConsensus, BaseHead


@HEADS.register_module
Expand Down

0 comments on commit 5e84344

Please sign in to comment.