迁移学习损失函数MMD(最大均值化差异)–python代码实现

MMD介绍

MMD（Max mean discrepancy 最大均值差异）是迁移学习，尤其是Domain adaptation （域适应）中使用最广泛（目前）的一种损失函数，主要用来度量两个不同但相关的分布的距离。两个分布的距离定义为：
$$MMD(X,Y)={\Vert \frac{1}{n}\sum _{i=1}^n\varphi \left(x_i\right)-\frac{1}{m}\sum _{j=1}^m\varphi \left(y_j\right)\Vert }_H^2$$

主代码编写

该代码基于torch.version = ‘1.9.0’

import torch
import torch.nn as nn

class MMDLoss(nn.Module):
    '''
    计算源域数据和目标域数据的MMD距离
    Params:
    source: 源域数据（n * len(x))
    target: 目标域数据（m * len(y))
    kernel_mul:
    kernel_num: 取不同高斯核的数量
    fix_sigma: 不同高斯核的sigma值
    Return:
    loss: MMD loss
    '''
    def __init__(self, kernel_type='rbf', kernel_mul=2.0, kernel_num=5, fix_sigma=None, **kwargs):
        super(MMDLoss, self).__init__()
        self.kernel_num = kernel_num
        self.kernel_mul = kernel_mul
        self.fix_sigma = None
        self.kernel_type = kernel_type

    def guassian_kernel(self, source, target, kernel_mul, kernel_num, fix_sigma):
        n_samples = int(source.size()[0]) + int(target.size()[0])
        total = torch.cat([source, target], dim=0)
        total0 = total.unsqueeze(0).expand(
            int(total.size(0)), int(total.size(0)), int(total.size(1)))
        total1 = total.unsqueeze(1).expand(
            int(total.size(0)), int(total.size(0)), int(total.size(1)))
        L2_distance = ((total0-total1)**2).sum(2)
        if fix_sigma:
            bandwidth = fix_sigma
        else:
            bandwidth = torch.sum(L2_distance.data) / (n_samples**2-n_samples)
        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)

    def linear_mmd2(self, f_of_X, f_of_Y):
        loss = 0.0
        delta = f_of_X.float().mean(0) - f_of_Y.float().mean(0)
        loss = delta.dot(delta.T)
        return loss

    def forward(self, source, target):
        if self.kernel_type == 'linear':
            return self.linear_mmd2(source, target)
        elif self.kernel_type == 'rbf':
            batch_size = int(source.size()[0])
            kernels = self.guassian_kernel(
                source, target, kernel_mul=self.kernel_mul, kernel_num=self.kernel_num, fix_sigma=self.fix_sigma)
            XX = torch.mean(kernels[:batch_size, :batch_size])
            YY = torch.mean(kernels[batch_size:, batch_size:])
            XY = torch.mean(kernels[:batch_size, batch_size:])
            YX = torch.mean(kernels[batch_size:, :batch_size])
            loss = torch.mean(XX + YY - XY - YX)
            return loss

程序验证

##在这里 第2维一定要相同，否则报错
source = torch.rand(64,14)  # 可以理解为源域有64个14维数据
target = torch.rand(32,14)  # 可以理解为源域有32个14维数据
print(target)

>>>output
tensor([[0.9035, 0.0088, 0.5867, 0.5595, 0.9350, 0.2739, 0.8775, 0.5562, 0.5402,
         0.5242, 0.4745, 0.7307, 0.7791, 0.7420],
        [0.2798, 0.6476, 0.3744, 0.5406, 0.3941, 0.6669, 0.2026, 0.8296, 0.3071,
         0.9042, 0.4810, 0.5235, 0.0547, 0.9110],
        [0.8051, 0.0702, 0.7907, 0.9708, 0.5310, 0.5851, 0.7881, 0.9082, 0.5963,
         0.9400, 0.3670, 0.8042, 0.5024, 0.2368],
        [0.5021, 0.7290, 0.3521, 0.6293, 0.8796, 0.2098, 0.0304, 0.9125, 0.3285,
         0.8485, 0.6877, 0.5695, 0.9506, 0.0752],
        [0.0798, 0.7908, 0.2785, 0.1369, 0.6762, 0.3342, 0.4930, 0.1807, 0.5963,
         0.2114, 0.4937, 0.4692, 0.3694, 0.9456],
         ...
        [0.1638, 0.7100, 0.9024, 0.5154, 0.8746, 0.8611, 0.1314, 0.0308, 0.6660,
         0.3719, 0.6827, 0.6789, 0.2416, 0.4617],
        [0.4449, 0.8304, 0.4036, 0.0563, 0.3832, 0.3553, 0.7947, 0.9335, 0.2704,
         0.9798, 0.2621, 0.4497, 0.9440, 0.7362]])

MMD = MMDLoss()
a = MMD(source=source, target=target)
print(a)

>>>output
tensor(0.1448)

嵌入到CNN中代码实现

先定义一个简单的CNN模型

class Net_only(nn.Module):
    '''
    计算源域数据和目标域数据的MMD距离
    Params:
    x_in: 输入数据（batch, channel, hight, width）
    Return:
    x_out: 输出数据（batch, n_labes)
    '''
    ## 这里 x_in：batch=64, channel=3, hight=128, width=128
    ## x_out：batch=64, n_labes=5
    def __init__(self):
        super(Net_only, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, 3)
        self.pool = nn.MaxPool2d(2, 2)
        self.bn1 = nn.BatchNorm2d(32)
        self.conv2 = nn.Conv2d(32, 64, 3)
        self.pool = nn.MaxPool2d(2, 2)
        self.bn2 = nn.BatchNorm2d(64)
        self.conv3 = nn.Conv2d(64, 64, 3)
        self.pool = nn.MaxPool2d(2, 2)
        self.bn3 = nn.BatchNorm2d(64)
        self.conv3 = nn.Conv2d(64, 64, 3)
        self.drop1d = nn.Dropout(0.2)
        self.bn4 = nn.BatchNorm2d(64)
        self.fc1 = nn.Linear(64 * 14 * 14, 1024)
        self.fc2 = nn.Linear(1024, 256)
        self.fc3 = nn.Linear(256, 5)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.bn1(x)
        x = self.pool(F.relu(self.conv2(x)))
        x = self.bn2(x)
        x = self.pool(F.relu(self.conv3(x)))
        x = self.bn3(x)
        x = x.view(-1, x.size(1) * x.size(2) * x.size(3))
        x = F.relu(self.fc1(x))
        x = self.drop1d(x)
        x = F.relu(self.fc2(x))
        x = self.drop1d(x)
        x = self.fc3(x)
        return x

对CNN模型进行测试

model = Net_only()
source = torch.rand(64, 3, 128, 128) # 模拟产生batch=64，channel=3， hight=128， width=128 的源域图片数据
target = torch.rand(32, 3, 128, 128) # 模拟产生batch=32，channel=3， hight=128， width=128 的源域图片数据
source = model(source)
target = model(target)
print(source.shape)

>>>output
torch.Size([64, 5])

现在计算MMD损失

MMD = MMDLoss()
loss = MMD(source=source, target=target)
print(loss)

>>>output
tensor(0.0884, grad_fn=<MeanBackward0>)

迁移学习损失的运用

loss = clf_loss + lamb * transfer_loss
clf_loss是源域的分类损失，transfer_loss即本篇所介绍的MMD_loss，lamb是超参数

总结

迁移损失MMD其输入X， Y分别是souce = Net(source)，target = Net(target)，也就是模型的输出。
参考资料：
链接: 王晋东github
链接:https://blog.csdn.net/a529975125/article/details/81176029
欢迎关注公众号：故障诊断与python学习