loss_builder.py

Go to the documentation of this file.

from smrunner.losses.classification_loss import CrossEntropyLoss, MutiTaskCrossEntropyLoss, MutiTaskSigmoidFocalLoss, MutiTaskMultiLabelCrossEntropyLoss
 
 
__loss_builder__ = {
    "CrossEntropyLoss" : CrossEntropyLoss,
    "MutiTaskCrossEntropyLoss" : MutiTaskCrossEntropyLoss,
    "MutiTaskSigmoidFocalLoss" : MutiTaskSigmoidFocalLoss,
    "MutiTaskMultiLabelCrossEntropyLoss" : MutiTaskMultiLabelCrossEntropyLoss
 
}
 
class CombinedLossEvaluators(object):
    """
    다양한 로스의 결합을 지원하는 클래스
    args:
        loss_evaluators (list): 로스 클래스 인스턴스 리스트
        loss_weights (list[float]): 각 로스별 가중치
    """
    def __init__(self, loss_evaluators, loss_weights):        
        self.loss_evaluators = loss_evaluators
        self.loss_weights = loss_weights
 
    def __call__(self, data):        
        losses = []
        for weight, loss_evaluator in zip(self.loss_weights, self.loss_evaluators):
            losses.append(weight*loss_evaluator(data))
        
        return sum(losses)
 
 
def build_loss(cfg):
    """
    모델의 학습에 사용할 로스 계산기를 빌드하는 기능
    args:
        cfg (str or Config): 학습 파라미터가 정의된 모델의 config 파일 경로 또는 mmengine.config.Config.fromfile()로 생성한 객체
    """
 
    #config 로드
    if isinstance(cfg, str):
        cfg = Config.fromfile(cfg_path)
 
    loss_evaluators = []
    loss_weights = []
 
    losses = cfg.loss #로스 파라미터
    loss_types = losses.keys() #로스 클래스 목록
 
    for type in loss_types:
 
        #예외처리
        assert type in __loss_builder__, \
            f'not found loss type : {type}'
 
        loss_cfg = losses.get(type) #losses[type]
        loss_weights.append(loss_cfg.weight) #로스에 대한 가중치
        loss_cfg.pop('weight')
        loss_evaluators.append(__loss_builder__[type](**loss_cfg)) #로스 클래스 빌드
        
    #로스의 결합 클래스 빌드
    return CombinedLossEvaluators(loss_evaluators, loss_weights)