loss.py

import torch.nn.functional as F
import torch
from torch import einsum
class coxph_loss(torch.nn.Module):

    def __init__(self):
        super(coxph_loss, self).__init__()

    def forward(self, sorted_risk, censors):

        #riskmax = risk
        riskmax = F.normalize(sorted_risk, p=2, dim=0)

        log_risk = torch.log((torch.cumsum(torch.exp(riskmax), dim=0)))

        uncensored_likelihood = torch.add(riskmax, -log_risk)
        resize_censors = censors.resize_(uncensored_likelihood.size()[0], 1)
        censored_likelihood = torch.mul(uncensored_likelihood, resize_censors)

        loss = -torch.sum(censored_likelihood) / float(censors.nonzero().size(0))
        #loss = -torch.sum(censored_likelihood) / float(censors.size(0))

        return loss

class pos_coxph_loss(torch.nn.Module):

    def __init__(self):
        super(pos_coxph_loss, self).__init__()

    def forward(self, sorted_risk, censors):
        idx = torch.argsort(sorted_risk,descending=True,dim=0).squeeze(-1)
        sorted_risk = sorted_risk[idx]
        censors = censors[idx]

        #riskmax = risk
        riskmax = F.normalize(sorted_risk, p=2, dim=0)

        log_risk = torch.log((torch.cumsum(torch.exp(riskmax), dim=0)))

        uncensored_likelihood = torch.add(riskmax, -log_risk)
        resize_censors = censors.resize_(uncensored_likelihood.size()[0], 1)
        censored_likelihood = torch.mul(uncensored_likelihood, resize_censors)

        loss = -torch.sum(censored_likelihood) / float(censors.nonzero().size(0))
        #loss = -torch.sum(censored_likelihood) / float(censors.size(0))

        return loss

class mse_loss(torch.nn.Module):

    def __init__(self):
        
        super(mse_loss, self).__init__()
        self.loss = torch.nn.MSELoss()

    def forward(self, sorted_risk, sorted_gt, censors):
        censored_sorted_risk = sorted_risk.squeeze(-1)*censors
        censored_sorted_gt = sorted_gt*censors

        loss = self.loss(censored_sorted_risk,censored_sorted_gt)

        return loss
    
class coxph_with_mse_loss(torch.nn.Module):
    def __init__(self):
        super(coxph_with_mse_loss, self).__init__()
        self.mseloss = mse_loss()
        self.coxph_loss = coxph_loss()
    def forward(self, sorted_risk, sorted_gt, censors):
        mse_loss = self.mseloss(sorted_risk, sorted_gt, censors)
        coxph_loss = self.coxph_loss(sorted_risk, censors)
        return mse_loss + 0.5 * coxph_loss

class bcr_with_mse_loss(torch.nn.Module):
    def __init__(self):
        super(bcr_with_mse_loss, self).__init__()
        self.mseloss = mse_loss()
        self.w = torch.nn.Parameter(torch.ones(128)*0.01)

    def forward(self, sorted_risk, sorted_gt, censors, x):
        mse_loss = self.mseloss(sorted_risk, sorted_gt, censors)
        N,C = x.size()
        dis = abs(torch.stack([x - x[i].unsqueeze(0) for i in range(N)]))
        bcr_loss = einsum('nvc,c->nv',dis,self.w.to(dis.device))
        bcr_loss = torch.sigmoid(bcr_loss)
        ones = torch.ones_like(bcr_loss)
        for i in range(N):
            ones[i][i]=0
        bcr_loss = -torch.log(bcr_loss) * ones
        bcr_loss = bcr_loss.sum()/(N * (N-1))
        return mse_loss + bcr_loss

class Regularization(object):
    def __init__(self, order, weight_decay):
        ''' The initialization of Regularization class

        :param order: (int) norm order number
        :param weight_decay: (float) weight decay rate
        '''
        super(Regularization, self).__init__()
        self.order = order
        self.weight_decay = weight_decay

    def __call__(self, model):
        ''' Performs calculates regularization(self.order) loss for model.

        :param model: (torch.nn.Module object)
        :return reg_loss: (torch.Tensor) the regularization(self.order) loss
        '''
        reg_loss = 0
        for name, w in model.named_parameters():
            if 'weight' in name:
                reg_loss = reg_loss + torch.norm(w, p=self.order)
        reg_loss = self.weight_decay * reg_loss
        return reg_loss    
class NegativeLogLikelihood(torch.nn.Module):
    def __init__(self):
        super(NegativeLogLikelihood, self).__init__()
        self.L2_reg = 0# config['l2_reg']
        self.reg = Regularization(order=2, weight_decay=self.L2_reg)

    def forward(self, risk_pred, y, e, model):
        if len(y.shape) ==1:
           y = y.unsqueeze(-1)
        mask = torch.ones(y.shape[0], y.shape[0]).to(risk_pred.device)
        
        mask[(y.T - y) > 0] = 0
        log_loss = torch.exp(risk_pred) * mask
        log_loss = torch.sum(log_loss, dim=0) / torch.sum(mask, dim=0)
        log_loss = torch.log(log_loss).reshape(-1, 1)
        neg_log_loss = -torch.sum((risk_pred-log_loss) * e) / torch.sum(e)
        l2_loss = self.reg(model)
        return neg_log_loss # + l2_loss


def guassian_kernel(source, target, kernel_mul=2.0, kernel_num=5, fix_sigma=None):
    '''
    将源域数据和目标域数据转化为核矩阵,即上文中的K
    Params: 
	    source: 源域数据(n * len(x))
	    target: 目标域数据(m * len(y))
	    kernel_mul: 
	    kernel_num: 取不同高斯核的数量
	    fix_sigma: 不同高斯核的sigma值
	Return:
		sum(kernel_val): 多个核矩阵之和
    '''
    n_samples = int(source.size()[0])+int(target.size()[0])# 求矩阵的行数，一般source和target的尺度是一样的，这样便于计算
    total = torch.cat([source, target], dim=0)#将source,target按列方向合并
    #将total复制（n+m）份
    total0 = total.unsqueeze(0).expand(int(total.size(0)), int(total.size(0)), int(total.size(1)))
    #将total的每一行都复制成（n+m）行，即每个数据都扩展成（n+m）份
    total1 = total.unsqueeze(1).expand(int(total.size(0)), int(total.size(0)), int(total.size(1)))
    #求任意两个数据之间的和，得到的矩阵中坐标（i,j）代表total中第i行数据和第j行数据之间的l2 distance(i==j时为0）
    L2_distance = ((total0-total1)**2).sum(2) 
    #调整高斯核函数的sigma值
    if fix_sigma:
        bandwidth = fix_sigma
    else:
        bandwidth = torch.sum(L2_distance.data) / (n_samples**2-n_samples)
    #以fix_sigma为中值，以kernel_mul为倍数取kernel_num个bandwidth值（比如fix_sigma为1时，得到[0.25,0.5,1,2,4]
    bandwidth /= kernel_mul ** (kernel_num // 2)
    bandwidth_list = [bandwidth * (kernel_mul**i) for i in range(kernel_num)]
    #高斯核函数的数学表达式
    kernel_val = [torch.exp(-L2_distance / bandwidth_temp) for bandwidth_temp in bandwidth_list]
    #得到最终的核矩阵
    return sum(kernel_val)#/len(kernel_val)

def mmd_rbf(source, target, kernel_mul=2.0, kernel_num=5, fix_sigma=None):
    '''
    计算源域数据和目标域数据的MMD距离
    Params: 
	    source: 源域数据(n * len(x))
	    target: 目标域数据(m * len(y))
	    kernel_mul: 
	    kernel_num: 取不同高斯核的数量
	    fix_sigma: 不同高斯核的sigma值
	Return:
		loss: MMD loss
    '''
    batch_size = int(source.size()[0])#一般默认为源域和目标域的batchsize相同
    kernels = guassian_kernel(source, target,
        kernel_mul=kernel_mul, kernel_num=kernel_num, fix_sigma=fix_sigma)
    #根据式（3）将核矩阵分成4部分
    XX = kernels[:batch_size, :batch_size]
    YY = kernels[batch_size:, batch_size:]
    XY = kernels[:batch_size, batch_size:]
    YX = kernels[batch_size:, :batch_size]
    loss = torch.mean(XX + YY - XY -YX)
    return loss#因为一般都是n==m，所以L矩阵一般不加入计算