timm_backbone.TimmBackbone Class Reference

Inheritance diagram for timm_backbone.TimmBackbone:

Public Member Functions
	__init__ (self, model_name, input_key, pretrained=False)
	forward (self, sample)

Public Attributes
	input_key
	timm_model

Detailed Description

timm 패키지의 백본을 사용하기 위한 클래스
timm 모델의 head는 nn.Identity()로 치완함

Definition at line 5 of file timm_backbone.py.

Constructor & Destructor Documentation

__init__()

timm_backbone.TimmBackbone.__init__	(		self,
			model_name,
			input_key,
			pretrained = False )

args:
    model_name (str) : timm 패키지에서 지원하는 모델 이름
    input_key (str) : inference에 사용하는 입력데이터의 키값
    pretrained (bool) : timm 에서 제공하는 프리트레인 파라미터 사용 여부

Definition at line 11 of file timm_backbone.py.

    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 : 파라미터 로드하는 코드 작성 필요
        

Member Function Documentation

forward()

timm_backbone.TimmBackbone.forward	(		self,
			sample )

Definition at line 31 of file timm_backbone.py.

    def forward(self, sample):
        x = sample[self.input_key]
        return self.timm_model(x)

Member Data Documentation

input_key

timm_backbone.TimmBackbone.input_key

Definition at line 20 of file timm_backbone.py.

timm_model

timm_backbone.TimmBackbone.timm_model

Definition at line 23 of file timm_backbone.py.

---

The documentation for this class was generated from the following file:

• smaction/models/backbones/timm_backbone.py