classification_loss.py

Go to the documentation of this file.

import torch
import torch.nn as nn
from torch.nn import functional as F
 
class CrossEntropyLoss(nn.Module):
    """
    모델의 출력(라벨 or 테스크)이 여러개일때 크로스엔트로피 로스를 계산하기위한 클래스
    args:
        pred_keys (list[str]): 모델이 예측한 결과(inference) 값이 저장된 키
        gt_keys (list[str]): 학습 데이터가 저장된 키, pred_keys에 대응하는 키가 순서에 맞게 들어있어야함
        weights (list[float]): 각 테스크의 로스에 대한 가중치, 설정하지 않으면 동일 가중치 사용
        data_num (list[int] or None): 클래스의 데이터 수량
        device (str): 모델이 구동하는 디바이스
    return (Tensor): 각 테스크 로스의 가중합
    """
    def __init__(self, pred_keys, gt_keys, weights=None, data_num=None, device='cpu'):
        super(CrossEntropyLoss, self).__init__()
        
        
        self.pred_keys = pred_keys
        self.gt_keys = gt_keys
        self.weights = weights
        if weights is None:
            self.weights = [1.0]*len(pred_keys)
 
        self.loss = nn.ModuleDict()
        for key in pred_keys:
 
            if data_num is None:
                self.loss[key] = nn.CrossEntropyLoss()
            else:        
                d_num = data_num[key]
                s_num = sum(d_num)
                weight = [n/s_num for n in d_num]
                weight = torch.tensor(weight).to(device)
                self.loss[key] = nn.CrossEntropyLoss(weight=weight)
 
    def calc_loss(self, pred, gt, pred_key):
        return self.loss[pred_key](pred, gt)
    
 
    def forward(self, data):
        losses = []
 
        for pred_key, gt_key, weight in zip(self.pred_keys, self.gt_keys, self.weights):
            losses.append( weight * self.calc_loss(data[pred_key], data[gt_key], pred_key) )
 
        return sum(losses)
 
 
class MutiTaskCrossEntropyLoss(nn.Module):
    """
    멀티 라벨 or 멀티 테스크 문제를 학습하기 위한 크로스엔트로피 클래스
    멀티 라벨 데이터를 기반, 멀티 라벨 각각을 테스크로 정의함, 학습데이터는 타겟 테스크에 대한 라벨값을 가지고 있음
    학습 데이터는 타겟 테스크가 설정되어 있음, 하나의 학습 데이터는 타겟 테스크에 대한 경로만 학습함
    args:
        task_key (str): 학습할 테스크 또는 라벨에 대한 정보가 설정된 키
        pred_keys (list[str]): 모델이 예측한 결과(inference) 값이 저장된 키
        gt_keys (list[str]): 학습 데이터가 저장된 키, pred_keys에 대응하는 키가 순서에 맞게 들어있어야함
        target_tasks (list[str]): 학습 데이터의 타겟 테스크, task_key에 설정됨
        weights (list[float]): 각 테스크의 로스에 대한 가중치, 설정하지 않으면 동일 가중치 사용
        data_num (dict or None): 테스크별 클래스의 데이터 수량
        device (str): 모델이 구동하는 디바이스
    return (Tensor): 각 테스크별 로스의 가중합
    """
    def __init__(self, task_key, pred_keys, gt_keys, target_tasks,  weights=[1.0], data_num=None, device='cpu'):
        super(MutiTaskCrossEntropyLoss, self).__init__()
        self.task_key = task_key
        self.target_tasks = target_tasks
        self.weights = weights
        if weights is None:
            self.weights = [1.0]*len(pred_keys)
        self.pred_keys = pred_keys
        self.gt_keys = gt_keys
 
        self.loss = nn.ModuleDict()
        for key in pred_keys:
 
            if data_num is None:
                self.loss[key] = nn.CrossEntropyLoss(ignore_index=-1)
            else:        
                d_num = data_num[key]
                s_num = sum(d_num)
                weight = [n/s_num for n in d_num]
                weight = torch.tensor(weight).to(device)
                self.loss[key] = nn.CrossEntropyLoss(weight=weight, ignore_index=-1)
 
    def calc_loss(self, pred, gt, pred_key):
        return self.loss[pred_key](pred, gt)
    
 
    def forward(self, data):
        losses = []
 
        for pred_key, gt_key, target_task, weight in zip(self.pred_keys, self.gt_keys, self.target_tasks, self.weights):
            
            #task에 해당하는 샘플 인덱스
            target_idx = [task == target_task for task in data[self.task_key]]
            if sum(target_idx) == 0:
                continue
                
            target_idx = torch.tensor(target_idx)
 
            #로스 계산
            losses.append( weight * self.calc_loss(data[pred_key][target_idx], data[gt_key][target_idx], pred_key) )
 
        return sum(losses)
 
class MutiTaskMultiLabelCrossEntropyLoss(nn.Module):
    """
    멀티 라벨, 멀티 테스크 문제를 학습하기 위한 크로스엔트로피 클래스
    멀티 라벨 데이터를 기반, 멀티 라벨 각각을 테스크로 정의함, 학습데이터는 모든 테스크에 대한 라벨값을 가지고 있음(멀티 라벨 데이터)
    학습 데이터는 메인 테스크가 설정되어 있음, 하나의 학습 데이터에 대해 모든 테스크에 대해 로스를 계산하고 가중합을 취함
    메인 테스크와 메인이 아닌 테스크에 대한 중요도 조절을 위해 가중치 조정이 필요함
    TODO
        모든 라벨에 대해서 학습하기 때문에 데이터 불균현 문제가 존재함 -> 해결 방안이 필요함
    args:
        task_key (str): 학습할 테스크 또는 라벨에 대한 정보가 설정된 키
        pred_keys (list[str]): 모델이 예측한 결과(inference) 값이 저장된 키
        gt_keys (list[str]): 학습 데이터가 저장된 키, pred_keys에 대응하는 키가 순서에 맞게 들어있어야함
        target_tasks (list[str]): 학습 데이터의 메인 테스크, task_key에 설정됨
        train_tasks (list[str]): 학습할 모든 테스크(라벨)
        target_task_weight (float): 메인 테스크에 대한 가중치
        non_target_task_weight (float): 메인이 아닌 테스크에 대한 가중치
        weights (list[float]): 각 테스크의 로스에 대한 가중치, 설정하지 않으면 동일 가중치 사용
    return (Tensor): 각 테스크별 로스의 가중합
    """
    def __init__(self, task_key, pred_keys, gt_keys, target_tasks, train_tasks, target_task_weight=1.0, non_target_task_weight=0.025, weights=None, **args):
        super(MutiTaskMultiLabelCrossEntropyLoss, self).__init__()
        self.task_key = task_key
        self.pred_keys = pred_keys
        self.gt_keys = gt_keys
        self.target_tasks = target_tasks
        self.train_tasks = train_tasks
        self.weights = weights
        if weights is None:
            self.weights = [1.0]*len(pred_keys)
        
        self.target_task_weight = target_task_weight
        self.non_target_task_weight = non_target_task_weight
 
        self.loss = nn.CrossEntropyLoss(ignore_index=-1)
        
 
    def calc_loss(self, pred, gt):
        return self.loss(pred, gt)
    
 
    def forward(self, data):
        
        losses = []
 
        for pred_key, gt_key, target_task, weight in zip(self.pred_keys, self.gt_keys, self.target_tasks, self.weights):
 
            #TODO : 코드가 비효율적으로 작성된것 처럼 보임, 효율적인 구조로 변경 필요
            if self.non_target_task_weight > 0: #하나의 샘플이 모든 테스크에 대해 학습
                for train_task in self.train_tasks: #학습 테스크
                    
                    #가중치 설정, 테스크 기본 가중치* 메인 or Non-Main 테크스 가중치
                    if train_task == target_task: #메인 테스크
                        weight = weight*self.target_task_weight
                    else: #메인 테스크가 아닌 테스크
                        weight = weight*self.non_target_task_weight
 
                    #학습 테스크 샘플 인덱스
                    idx = [ task == train_task for task in data[self.task_key]]
                    if sum(idx) == 0:
                        continue
                    
                    idx = torch.tensor(idx)
                    
                    #학습 테스크의 클래스가 -1만 있을경우 처리, -1은 크로스 엔트로피에 ignore_index로 설정되어 있음, 따라서 -1인 경우 학습되지 않음
                    if (data[gt_key][idx] != -1).sum() == 0:
                        continue
 
                    losses.append(weight*self.calc_loss(data[pred_key][idx], data[gt_key][idx]))
            else: #하나의 샘플이 메인 테스크에 대해서만 학습, MutiTaskCrossEntropyLoss와 동일
                #task에 해당하는 샘플 인덱스
                target_idx = [task == target_task for task in data[self.task_key]]
                if sum(target_idx) == 0:
                    continue
                    
                target_idx = torch.tensor(target_idx)
 
                #로스 계산
                losses.append( weight * self.calc_loss(data[pred_key][target_idx], data[gt_key][target_idx]) )
 
 
        return sum(losses)
 
 
class MutiTaskSigmoidFocalLoss(nn.Module):
    """
    멀티 라벨 또는 멀티 테스크 문제를 학습하기 위한 시그모이드 포컬로스 클래스
    멀티 라벨 데이터를 기반, 멀티 라벨 각각을 테스크로 정의함, 학습데이터는 타겟 테스크에 대한 라벨값을 가지고 있음
    학습 데이터는 타겟 테스크가 설정되어 있음, 하나의 학습 데이터는 타겟 테스크에 대한 경로만 학습함
    args:
        task_key (str): 학습할 테스크 또는 라벨에 대한 정보가 설정된 키
        pred_keys (list[str]): 모델이 예측한 결과(inference) 값이 저장된 키
        gt_keys (list[str]): 학습 데이터가 저장된 키, pred_keys에 대응하는 키가 순서에 맞게 들어있어야함
        target_tasks (list[str]): 학습 데이터의 테스크, task_key에 설정됨, 해당 테스크 경로만 학습함
        weights (list[float]): 각 테스크의 로스에 대한 가중치, 설정하지 않으면 동일 가중치 사용
        alpha (float): 클래스 불균형 파라미터
        gamma (float): 샘플 난이도 가중치
        beta (float): 테스크 난이도에 대한 가중치
            TODO: pjm 추가, 성능이 좋지 않다고 판단되면 제거
        reduction (str): 최종 출력 로스에 적용할 연산
            'mean': 평균
            'sum': 합
            None : 연산이 적용되지 않음
    return (Tensor): 각 테스크별 로스의 가중합
    """
    def __init__(self, task_key, pred_keys, gt_keys, target_tasks,  weights=None, alpha=0.25, gamma=2, beta=0, reduction='mean', **args):
        super(MutiTaskSigmoidFocalLoss, self).__init__()
 
        self.task_key = task_key
        self.pred_keys = pred_keys
        self.gt_keys = gt_keys
        self.target_tasks = target_tasks
        self.weights = weights
        if weights is None:
            self.weights = [1.0] * len(pred_keys)
            
        self.alpha = alpha
        self.gamma = gamma
        self.beta = beta
        self.reduction = reduction
 
    def sigmoid_focal_loss(self, inputs, targets):
        """
        시그모이드 포컬로스 계산하는 기능
        args:
            inputs (Tensor): 모델의 출력 결과
            targets (Tensor): 모델의 출력에 대응하는 GT
        return (Tensor)
        """
 
        inputs = inputs.float() #타입변환
        targets = targets.float() #타입변환
 
        #시그모이드 포컬로스 계산
        p = torch.sigmoid(inputs) #시그모이드
        ce_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction="none") #바이너리 크로스엔트로피
        p_t = p * targets + (1 - p) * (1 - targets)
        loss = ce_loss * ((1 - p_t) ** self.gamma) #샘플 난이도 가중치
 
        if self.alpha >= 0:
            #클래스 불균형 조정
            alpha_t = self.alpha * targets + (1 - self.alpha) * (1 - targets)
            loss = alpha_t * loss
 
        if self.beta > 0:
            #테스크 난이도 가중치
            w = torch.sum((1 - p_t) * targets, dim=1).mean() ** self.beta
            #w = (1 - p_t).mean() ** self.beta
            loss = w * loss
 
        #reduction
        if self.reduction == "mean":
            loss = loss.mean()
        elif self.reduction == "sum":
            loss = loss.sum()
 
        return loss
 
    def calc_loss(self, pred, gt):
        """
        시그모이드 포컬로스를 계산하는 기능
        바이너리 크로스엔트로피 적용을 위해 GT에 대한 전처리(one hot encording) 장치가 추가됨
        args:
            pred (Tensor): 모델의 예측 결과(logit)
            gt (Tensor): 클래스 번호
        """
        class_num = pred.shape[1] #클래스 수
        gt_one_hot = F.one_hot(gt, class_num) #one hot encording
 
        #로스 계산
        return self.sigmoid_focal_loss(pred, gt_one_hot)
    
 
    def forward(self, data):
        losses = []
 
        for pred_key, gt_key, target_task, weight in zip(self.pred_keys, self.gt_keys, self.target_tasks, self.weights):
            
            #task에 해당하는 샘플 인덱스
            target_idx = [task == target_task for task in data[self.task_key]]
            if sum(target_idx) == 0:
                continue
            
            target_idx = torch.tensor(target_idx)
 
            #로스 계산
            losses.append( weight * self.calc_loss(data[pred_key][target_idx], data[gt_key][target_idx]) )
 
        return sum(losses)