timm_backbone.py

Go to the documentation of this file.

import timm
import torch
import torch.nn as nn
 
class TimmBackbone(nn.Module):
    """
    timm 패키지의 백본을 사용하기 위한 클래스
    timm 모델의 head는 nn.Identity()로 치완함
    """
 
    def __init__(self, model_name, input_key, pretrained=False):
        """
        args:
            model_name (str) : timm 패키지에서 지원하는 모델 이름
            input_key (str) : inference에 사용하는 입력데이터의 키값
            pretrained (bool) : timm 에서 제공하는 프리트레인 파라미터 사용 여부
        """
        super().__init__()
 
        self.input_key = input_key
 
        if isinstance(pretrained, bool):
            self.timm_model = timm.create_model(model_name, pretrained=pretrained)
            self.timm_model.head = nn.Identity()
        else:
            self.timm_model = timm.create_model(model_name, pretrained=False)
            self.timm_model.head = nn.Identity()
            
            #TODO : 파라미터 로드하는 코드 작성 필요
        
    def forward(self, sample):
        x = sample[self.input_key]
        return self.timm_model(x)