mlp_head.py

Go to the documentation of this file.

import torch
import torch.nn as nn
 
class MLPHead(nn.Module):
    """
    여러 텐서들을 입력으로 다층 레이어를 기반으로 클래스 스코어를 출력하는 기능
    3차원 또는 1차원 입력 텐서들을 받아서 1차원 테서로 풀링한 뒤 하나의 텐서로 이어 붙임
    이후 MLP레이어를 통과시켜 클래스 스코어를 출력함
    args:
        in_channels (int) : 입력 텐서들의 채널 합
        num_classes (int) : 클래스 수
        layer_channels (list[int]) : 각 측의 레이어 채널 수
        dropout_ratio (int) : 드랍아웃 비율
        input_key (list[str]) : 모듈의 inference에 사용하는 입력데이터의 키값들
        input_type (str) : input_key에 대응하는 데이터의 구조
            ex) ['linear', '3d']
    """
    def __init__(self, in_channels, num_classes, layer_channels, dropout_ratio, input_key, input_type='linear'):
        super().__init__()
 
        self.dropout_ratio = dropout_ratio
        self.input_key = input_key
 
        layer_num = len(layer_channels)
        layer_list = []
 
        in_ch = in_channels
        for i in range(layer_num):
            out_ch = layer_channels[i]
            if self.dropout_ratio != 0:
                layer_list.append(nn.Dropout(p=self.dropout_ratio))
            layer_list.append(nn.Linear(in_ch, out_ch))
            layer_list.append(nn.BatchNorm1d(out_ch))
            layer_list.append(nn.ReLU(inplace=True))
            in_ch = out_ch
        
        if self.dropout_ratio != 0:
            layer_list.append(nn.Dropout(p=self.dropout_ratio))
        
        layer_list.append(nn.Linear(in_ch, num_classes))
 
        self.mlp = nn.Sequential(*layer_list)
 
        if input_type == '3d':
            self.avg_pool = nn.AdaptiveAvgPool3d((1, 1, 1))
        else:
            self.avg_pool = None
        
 
    def forward(self, feat_dict):
        """
        args:
            sample (dict)) : 입력 데이터, self.input_key에 해당하는 키가 있어야함
                self.input_key의 아이템은 Tensor 타입 -> shape (B, C), (B, C, T, H, W)
                    B : 배치 크기
                    C : 채널 크기
                    T : 시간축 크기
 
        return (Tensor):
            1차원 특징 벡터 -> shape (B, num_classes)
        """
        x = feat_dict[self.input_key]
        if self.avg_pool is not None:
            x = self.avg_pool(x)
            x = x.view(x.shape[0], -1)
 
        return self.mlp(x)