i3d_fusion.py

Go to the documentation of this file.

import torch
import torch.nn as nn
 
class I3DFusion(nn.Module):
    """
    3차원 입력 텐서들을 연결하는 기능의 모듈
    입력 텐서들을 풀링하여 하나의 텐서로 이어붙이고 1차원으로 변환함, 추가적으로 fc+bn+relu 연산을 적용할 수 있음
    args:
        in_channels (int) : 입력 텐서 채널들의 합
        out_channels (int) : 출력 텐서의 채널 수, 0 이상값을 설정하면 fc+bn+relu 연산을 적용함
        dropout_ratio (float) : fc+bn+relu 연산을 사용할 경우 드랍아웃 비율을 설정하는 파라미터
        input_key (str) : inference에 사용하는 입력데이터의 키값들
    """
    def __init__(self, in_channels, out_channels, dropout_ratio, input_key):
        super().__init__()
        self.out_channels = out_channels
        self.dropout_ratio = dropout_ratio
        self.input_key = input_key
        self.avg_pool = nn.AdaptiveAvgPool3d((1, 1, 1))
 
        if self.out_channels > 0:
            self.fc = nn.Linear(in_channels, out_channels)
            self.bn = nn.BatchNorm1d(out_channels)
            self.relu = nn.ReLU(inplace=True)
 
            if self.dropout_ratio != 0:
                self.dropout = nn.Dropout(p=self.dropout_ratio)
            else:
                self.dropout = None
 
    def forward(self, feat_dict):
        """
        args:
            sample (dict)) : 입력 데이터, self.input_key에 해당하는 키가 있어야함
                self.input_key의 아이템은 Tensor 타입 -> shape (B, C, T, H, W)
                    B : 배치 크기
                    C : 입력 채널
                    T : 시간
 
        return (Tensor):
            1차원 특징 벡터 -> shape (B, C_o)
                B : 배치 크기
                C_o : 채널 크기, out_channels 또는 입력 채널의 합
        """
        feats = []
        for key in self.input_key:
            x = feat_dict[key]
            x = self.avg_pool(x)
            feats.append(x)
        
        fusion_feat = torch.cat(feats, dim=1)
        fusion_feat = fusion_feat.view(x.shape[0], -1)
 
        if self.out_channels > 0:
            if self.dropout is not None:
                fusion_feat = self.dropout(fusion_feat)
 
            fusion_feat = self.relu(self.bn(self.fc(fusion_feat)))
 
            return fusion_feat
        else:
            return fusion_feat