linear_head.py

Go to the documentation of this file.

import torch
import torch.nn as nn
 
class LinearHead(nn.Module):
    """
    1차원 텐서를 입력으로 클래스 스코어를 출력하는 기능의 모듈
    1차원 입력 텐서에 fc 연산을 적용해서 클래스 스코어를 출력함
    args:
        in_channels (int) : 입력 텐서의 채널 수
        num_classes (int) : 클래스 수, 출력 텐서의 채널 수
        dropout_ratio (float) : 드랍아웃 비율을 설정하는 파라미터
        input_key (str) : inference에 사용하는 입력데이터의 키값
    """
    def __init__(self, in_channels, num_classes, dropout_ratio, input_key):
        super().__init__()
 
        self.dropout_ratio = dropout_ratio
        self.input_key = input_key
 
        self.fc_cls = nn.Linear(in_channels, num_classes)
        
        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, in_channels)
                    B : 배치 크기
 
        return (Tensor):
            1차원 특징 벡터 -> shape (B, num_classes)
        """
        x = feat_dict[self.input_key]
 
        if self.dropout is not None:
            x = self.dropout(x)
 
        cls_score = self.fc_cls(x)
 
        return cls_score