gcn_head.py

Go to the documentation of this file.

import torch.nn as nn
 
class GCNHead(nn.Module):
 
    def __init__(self,
                 in_channels,
                 num_class,
                 input_key):
        super().__init__()
 
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Linear(in_channels, num_class)
        self.input_key = input_key
        
 
    def forward(self, feat_dict):
        x = feat_dict[self.input_key]
        
        # N, M, C, T, V = x.shape
        N, C, T, V = x.shape
 
        # x = x.view(N * M, C, T, V)
 
        # global pooling
        x = self.pool(x)
        x = x.view(N, C)
        # x = x.mean(dim=1)
 
        cls_scores = self.fc(x)
 
        return cls_scores