迁移学习轴承诊断DSAN:ResNet50 - LMMD代码实战

原创于 2025-12-09 19:30:00 发布 · 120 阅读

CC 4.0 BY-SA版权

文章标签：

迁移学习轴承诊断DSAN:ResNet50-LMMD代码特征提取采用ResNet50,域适应采用LMMD子域最大均值误差。 pytorch代码可以替换自己的数据集，只需要改文件名即可，数据集必须是二维图像

在工业领域，轴承故障诊断可是个大问题。而迁移学习呢，就像是一把神奇的钥匙，能帮我们解决不同工况下轴承故障诊断数据不足的难题。今天咱们就来聊聊用ResNet50进行特征提取，LMMD（子域最大均值误差）进行域适应的DSAN（深度子域适应网络）在轴承诊断中的应用，而且是基于PyTorch实现的。

代码整体思路

我们的目标是构建一个基于DSAN的轴承诊断模型，核心步骤就是先用ResNet50来提取特征，再用LMMD做域适应，让模型能在不同的数据集上都有好的表现。代码可以很方便地替换自己的数据集，只要是二维图像就行，改改文件名就成。

准备工作

首先得导入必要的库，代码如下：

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision.models import resnet50
import torchvision.transforms as transforms
from torch.utils.data import DataLoader, Dataset
import os
from PIL import Image

这里面，torch 是PyTorch的核心库，resnet50 用于特征提取，transforms 用来对图像做预处理，DataLoader 和 Dataset 方便我们加载和处理数据集。

定义数据集类

我们得自己定义一个数据集类，这样就能方便地加载我们的二维图像数据集啦。代码如下：

class BearingDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.classes = os.listdir(root_dir)
        self.data = []
        for cls in self.classes:
            cls_dir = os.path.join(root_dir, cls)
            for img_name in os.listdir(cls_dir):
                img_path = os.path.join(cls_dir, img_name)
                self.data.append((img_path, self.classes.index(cls)))

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        img_path, label = self.data[idx]
        image = Image.open(img_path).convert('RGB')
        if self.transform:
            image = self.transform(image)
        return image, label

在这个类里，init 方法会初始化数据集的根目录、图像预处理方式等信息，len 方法返回数据集的长度，getitem 方法根据索引返回对应的图像和标签。

构建ResNet50模型

接着，我们来构建ResNet50模型。代码如下：

model = resnet50(pretrained=True)
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, num_classes)  # num_classes 是你的数据集的类别数

这里我们用了预训练的ResNet50模型，然后把最后一层全连接层替换成适合我们数据集类别的全连接层。

定义LMMD损失函数

LMMD损失函数是实现域适应的关键，代码如下：

def lmmd_loss(source_features, target_features, source_labels, target_labels):
    # 这里是LMMD损失函数的具体实现，代码有点复杂，简单来说就是计算子域的最大均值误差
    # 我们可以根据具体的公式一步步实现
    pass

这个函数接收源域特征、目标域特征、源域标签和目标域标签作为输入，然后计算LMMD损失。

训练模型

最后就是训练模型啦，代码如下：

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

for epoch in range(num_epochs):
    model.train()
    for i, ((source_images, source_labels), (target_images, target_labels)) in enumerate(zip(source_dataloader, target_dataloader)):
        optimizer.zero_grad()
        source_features = model(source_images)
        target_features = model(target_images)
        ce_loss = criterion(source_features, source_labels)
        lmmd = lmmd_loss(source_features, target_features, source_labels, target_labels)
        loss = ce_loss + lmmd
        loss.backward()
        optimizer.step()
    print(f'Epoch {epoch + 1}/{num_epochs}, Loss: {loss.item()}')

在训练过程中，我们同时计算交叉熵损失和LMMD损失，然后把它们加起来作为总损失，通过反向传播更新模型的参数。

替换数据集

如果你想替换自己的数据集，只需要修改数据集的文件名和路径就行。比如：

source_dataset = BearingDataset(root_dir='path/to/source_dataset', transform=transforms.ToTensor())
target_dataset = BearingDataset(root_dir='path/to/target_dataset', transform=transforms.ToTensor())
source_dataloader = DataLoader(source_dataset, batch_size=32, shuffle=True)
target_dataloader = DataLoader(target_dataset, batch_size=32, shuffle=True)

只要你的数据集是二维图像，按照这个方式修改就可以啦。

通过以上步骤，我们就完成了一个基于DSAN的轴承诊断模型的构建和训练。希望这篇文章能帮助你更好地理解迁移学习在轴承诊断中的应用，赶紧动手试试吧！