import torch
import torch.nn as nn


def dm01():
    # 真实分布
    P = torch.tensor([[0.7, 0.2, 0.1]], dtype=torch.float)
    # 模型分布
    Q = torch.tensor([[0.5, 0.3, 0.2]], dtype=torch.float)
    # 定义KL散度损失，则聚合方式应该是sum求和
    criterion = nn.KLDivLoss(reduction='sum')
    # 注意：如果希望手算和PyTorch的结果对应上
    # # 则input应该是：【log(Q)，target是P】，即：【预测值，真实值】
    kl_loss_class = criterion(torch.log(Q), P)
    print("kl_loss_class-->", kl_loss_class)


def dm02():
    # 创建KL散度对象，reduction=batchmean是PyTorch推荐使用方式，一般也使用这种方式
    kl_loss = nn.KLDivLoss(reduction='batchmean')

    # 学生模型需要是log_softmax，教室模型需要softmax
    input_log_probs = torch.log_softmax(torch.tensor([[0.7, 0.2, 0.1]]), dim=1)
    target_probs = torch.softmax(torch.tensor([[0.5, 0.3, 0.2]]), dim=1)

    loss = kl_loss(input_log_probs, target_probs)
    print("loss-->", loss)



def dm03():
    # 创建KL散度对象，reduction=batchmean是推荐使用方式
    kl_loss = nn.KLDivLoss(reduction='batchmean')

    # 学生模型需要是log_softmax，教室模型需要softmax
    # 定义温度T，经验值：2或者4，这里我们依次试验【0.5, 1, 2, 4, 10, 100】
    T = 100
    input_log_probs = torch.log_softmax(torch.tensor([[0.7, 0.2, 0.1]]) / T, dim=1)
    print("input_log_probs-->", input_log_probs)
    target_probs = torch.softmax(torch.tensor([[0.5, 0.3, 0.2]]) / T, dim=1)
    print("target_probs-->", target_probs)
    loss = kl_loss(input_log_probs, target_probs)
    print("loss-->", loss)


if __name__ == '__main__':
    # dm01()
    # dm02()
    dm03()