utils.py

Go to the documentation of this file.

 
 
import torch
import torch.nn as nn
 
import math
 
from smutils.utils_data import remove_items
 
def load_model(model, path, ignore_key=None, device='cuda:0'):
    """
    모델의 파라미터를 로드하는 기능
    args:
        model : 파이토치 모델
        path (str): 파라미터 경로
        ignore_key (list[str]): 로드하지 않는 파라미터
        device (str): 모델이 구동될 디바이스
    """
 
    print('model path : ', path) #모델 파라미터 경로 출력
 
    state_dict = torch.load(path, map_location=device) #파일 읽기
 
    #로드할 모듈 셋팅
    module_names = model.get_module_names() #모듈 네임가져오기, runner에 구현되어 있어야함
    if ignore_key is not None:
        module_names = remove_items(module_names, ignore_key)
    
    #모델 파라미터 로드
    for name in module_names:
        #모듈 가져오기, 없으면 메시지 출력하고 다음으로 이동
        try:
            m = getattr(model, name) 
        except AttributeError:
            print(f'The model does not have {name} attribute.')
            print('Please check the model’s get_module_names() function.')
            continue
 
        #파라미터 로드
        if name in state_dict.keys():
            _dict = state_dict[name]
            print('load ', name)
 
            for i, k in zip(m.state_dict(), _dict):
                param = _dict[k]
                m.state_dict()[i].copy_(param)
        else:
            print(f'{name} is missing from the parameter file.')
 
def save_model(model, save_path):
    """
    모델의 파라미터를 저장하는 기능
    args:
        model : 파이토치 모델
        save_path (str): 파라미터 저장 경로
    """
 
    #로드할 모듈 셋팅
    checkpoint_model = {}    
    module_names = model.get_module_names() #모듈 네임가져오기, runner에 구현되어 있어야함
    for name in module_names:
        try :
            #모듈 네임 - 파라미터 구조의 딕셔너리 생성
            checkpoint_model[name] = getattr(model, name).state_dict()
        except AttributeError:
            print(f'The model does not have {name} attribute.')
            print('Please check the model’s get_module_names() function.')
 
    #저장
    torch.save(checkpoint_model, save_path)
 
def load_optimizer(optimizer, path):
    """
    옵티마티저 파라미터 로드 기능
    args:
        optimizer : 파이토치 옵티마이저
        path (str): 옵티마이저 파라미터 로드 경로
    """
    checkpoint = torch.load(path)
    optimizer.load_state_dict(checkpoint['optimizer'])
 
 
def save_optimizer(optimizer, save_path):
    """
    옵티마티저 파라미터 저장 기능
    args:
        optimizer : 파이토치 옵티마이저
        save_path (str): 옵티마이저 파라미터 저장 경로
    """
    checkpoint_opt = {'optimizer': optimizer.state_dict()}
    torch.save(checkpoint_opt, save_path)
 
 
def init_weights(model):
    """
    모델 파라미터 초기화 기능
    args:
        model : 파이토치 모델
    """
    for m in model.modules():
        if isinstance(m, nn.Conv2d): #conv 2d 초기화
            n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
            m.weight.data.normal_(0, math.sqrt(2. / n))
        elif isinstance(m, nn.Conv3d): #conv 3d 초기화
            n = m.kernel_size[0] * m.kernel_size[1] * m.kernel_size[2] * m.out_channels
            m.weight.data.normal_(0, math.sqrt(2. / n))
        elif isinstance(m, nn.BatchNorm2d): #배치놈 2d 초기화
            m.weight.data.fill_(1)
            m.bias.data.zero_()
        elif isinstance(m, nn.BatchNorm3d): #배치놈 3d 초기화
            m.weight.data.fill_(1)
            m.bias.data.zero_()
        elif isinstance(m, nn.Linear): #Linear 초기화, 0으로 초기화
            m.bias.data.zero_()
 
# def init_weights(model):
#     for m in model.modules():
#         if isinstance(m, nn.Conv2d):
#             # n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
#             # m.weight.data.normal_(0, math.sqrt(2. / n))
#             m.weight.data.normal_(0.0, 0.02)
#             if m.bias is not None:
#                 m.bias.data.fill_(0)
#         elif isinstance(m, nn.Conv3d):
#             n = m.kernel_size[0] * m.kernel_size[1] * m.kernel_size[2] * m.out_channels
#             m.weight.data.normal_(0, math.sqrt(2. / n))
#         elif isinstance(m, nn.BatchNorm2d):
#             m.weight.data.normal_(1.0, 0.02)
#             m.bias.data.fill_(0)
#             # m.weight.data.fill_(1)
#             # m.bias.data.zero_()
#         elif isinstance(m, nn.BatchNorm3d):
#             m.weight.data.fill_(1)
#             m.bias.data.zero_()
#         elif isinstance(m, nn.Linear):
#             m.weight.data.normal_(0.0, 0.02)
#             if m.bias is not None:
#                 m.bias.data.zero_()