cat_layer.py

Go to the documentation of this file.

import torch
import torch.nn as nn
 
class CatLayer(nn.Module):
    """
    입력 텐서들을 연결하는 기능의 모듈
    입력 텐서들을 풀링하여 1차원으로 만들고 하나의 텐서로 이어붙이는 기능을 수행함
    args:
        input_key (list[str]) : 모듈의 inference에 사용하는 입력데이터의 키값들
        input_type (list[str]) : input_key에 대응하는 데이터의 구조
            ex) ['2d', '3d']
    """
    def __init__(self, input_key, input_type):
        super().__init__()
        self.input_key = input_key
        self.input_type = input_type
 
        self.avg_pool_dict = nn.ModuleDict()
        for key, in_type in zip(self.input_key, self.input_type):
            if in_type == '3d':
                self.avg_pool_dict[key] = nn.AdaptiveAvgPool3d((1, 1, 1))
            elif in_type == '2d':
                self.avg_pool_dict[key] = nn.AdaptiveAvgPool2d((1, 1))
 
    def forward(self, feat_dict):
        """
        args:
            sample (dict)) : 입력 데이터, self.input_key에 해당하는 키가 있어야함
                self.input_key의 아이템은 Tensor 타입 -> shape (B, C, T, H, W) or (B, C, H(or T), W)
                    B : 배치 크기
                    C : 입력 채널
                    T : 시간
 
        return (Tensor):
            1차원 특징 벡터 -> shape (B, C_o)
                B : 배치 크기
                C_o : 채널 크기, 입력 채널의 합
        """
        feats = []
        for key in self.input_key:
            x = feat_dict[key]
            x = self.avg_pool_dict[key](x)
            x = x.view(x.shape[0], -1)
            feats.append(x)
        
        cat_feat = torch.cat(feats, dim=1)
        
        return cat_feat