├── README.md ├── model.py ├── test.py ├── loss.py ├── Dataset.py └── main.py /README.md: -------------------------------------------------------------------------------- 1 | # The-unsupervised-bearing-fault-diagnosis-method-based-on-the-dual-framework-Siamese-network 2 | Fault diagnosis code based on twin network 3 | -------------------------------------------------------------------------------- /model.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | class DeepSiameseNetwork(nn.Module): 3 | def __init__(self, in_channels,): 4 | super(DeepSiameseNetwork, self).__init__() 5 | 6 | 7 | self.shared_layer = nn.Sequential( 8 | nn.Conv2d(in_channels, 96, kernel_size=11, stride=4, padding=2), 9 | nn.ReLU(), 10 | nn.CrossMapLRN2d(size=5, alpha=0.0001, beta=0.75, k=1.0), 11 | nn.MaxPool2d(kernel_size=3, stride=2, padding=1), 12 | nn.Conv2d(96, 128, kernel_size=5, stride=1, padding=2), 13 | nn.ReLU(), 14 | nn.CrossMapLRN2d(size=5, alpha=0.0001, beta=0.75, k=1.0), 15 | nn.MaxPool2d(kernel_size=3, stride=2, padding=1), 16 | nn.Conv2d(128, 384, kernel_size=3, stride=1, padding=1), 17 | nn.ReLU(), 18 | nn.Conv2d(384, 192, kernel_size=3, stride=1, padding=1), 19 | nn.ReLU(), 20 | nn.Conv2d(192, 192, kernel_size=3, stride=1, padding=1), 21 | nn.ReLU(), 22 | nn.Conv2d(192, 128, kernel_size=3, stride=1, padding=1), 23 | nn.ReLU(), 24 | nn.MaxPool2d(kernel_size=3, stride=2, padding=1) 25 | ) 26 | self.flattened_size = 6272 27 | 28 | self.fc_layers = nn.Sequential( 29 | nn.Linear(self.flattened_size, 200), 30 | nn.ReLU(), 31 | nn.Linear(200, 200), 32 | nn.ReLU(), 33 | nn.Linear(200, 1), 34 | nn.ReLU() 35 | ) 36 | 37 | def forward(self, x1, x2): 38 | out1 = self.shared_layer(x1) 39 | out1 = out1.view(out1.size(0), -1) 40 | out2 = self.shared_layer(x2) 41 | out2 = out2.view(out2.size(0), -1) 42 | out1 = self.fc_layers(out1) 43 | out2 = self.fc_layers(out2) 44 | return out1, out2 45 | -------------------------------------------------------------------------------- /test.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as F 3 | from model import DeepSiameseNetwork 4 | from Dataset import TestSiameseDatasetA, TestSiameseDatasetB 5 | from torch.utils.data import DataLoader 6 | # Step 1: 加载模型 7 | in_channels = 3 # RGB 图像 8 | height, width = 227, 227 9 | model = DeepSiameseNetwork(in_channels) 10 | model.load_state_dict(torch.load(r'D:\故障诊断\是否故障\切比雪夫\best_model.pth')) 11 | model.eval() 12 | # Step 2: 准备数据 13 | # 假设有一对需要验证的图像：img1 和 img2 14 | img1_path = r'D:\故障诊断\是否故障\验证\108验证\外圈' 15 | img2_path = r'D:\故障诊断\是否故障\验证\sqv验证\正常' 16 | 17 | # 创建数据集对象 18 | def test_code(img1_path, img2_path, model, label): 19 | test_dataset_A = TestSiameseDatasetA(img1_path) 20 | test_dataset_B = TestSiameseDatasetB(img2_path) 21 | test_dataset_A = DataLoader(test_dataset_A, batch_size=1, shuffle=False) 22 | test_dataset_B = DataLoader(test_dataset_B, batch_size=1, shuffle=False) 23 | 24 | accuracy = test1(test_dataset_A, test_dataset_B, model, label) 25 | return accuracy 26 | 27 | 28 | def test1(test_loader_A, test_loader_B, model, true_label): 29 | model.eval() 30 | total = 0 31 | correct_nosame = 0 32 | correct_same = 0 33 | print(len(test_loader_A)) 34 | print(len(test_loader_B)) 35 | 36 | with torch.no_grad(): 37 | for data_A, data_B in zip(test_loader_A, test_loader_B): 38 | output_A, output_B = model(data_A, data_B) 39 | dist = F.pairwise_distance(output_A, output_B) 40 | target = torch.zeros_like(dist) # 假设目标为0 41 | predicted_euclidean = (dist < 0.5).float() 42 | correct_nosame += (predicted_euclidean == target).sum().item() 43 | 44 | # 计算与真实标签的比较结果 45 | correct_same += (predicted_euclidean == 1).sum().item() # 标签为1的数量即为正确数量 46 | 47 | total += target.size(0) 48 | 49 | # 计算准确率 50 | nosame = correct_nosame / total 51 | same = correct_same / total 52 | 53 | return nosame, same 54 | 55 | 56 | # 调用测试函数 57 | label_tensor = torch.tensor(1) # 将标签封装成张量 # 标签，1表示相似，0表示不相似 58 | 59 | accuracy = test_code(img1_path, img2_path, model, label_tensor) 60 | print("Accuracy第一个数是不相似，第二个数是相似：", accuracy) 61 | #第一个数是不相似，第二个数是相似 -------------------------------------------------------------------------------- /loss.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as F 3 | 4 | 5 | class ContrastiveLoss(torch.nn.Module): 6 | 7 | def __init__(self, margin): 8 | super(ContrastiveLoss, self).__init__() 9 | self.margin = margin 10 | 11 | def forward(self, x1, x2, y): 12 | dist = F.pairwise_distance(x1, x2) 13 | total_loss = (1 - y) * torch.pow(dist, 2) + \ 14 | y * torch.pow(torch.clamp_min_(self.margin - dist, 0), 2) 15 | loss = torch.mean(total_loss) 16 | return loss 17 | 18 | def manhattan_distance(x1, x2): 19 | return torch.sum(torch.abs(x1 - x2), dim=-1) 20 | class ContrastiveLoss_1(torch.nn.Module): 21 | 22 | def __init__(self, margin): 23 | super(ContrastiveLoss_1, self).__init__() 24 | self.margin = margin 25 | 26 | 27 | def forward(self, x1, x2, y): 28 | dist = manhattan_distance(x1, x2) 29 | total_loss = (1 - y) * torch.pow(dist, 2) + \ 30 | y * torch.pow(torch.clamp_min_(self.margin - dist, 0), 2) 31 | loss = torch.mean(total_loss) 32 | return loss 33 | 34 | 35 | class ContrastiveLoss_2(torch.nn.Module): 36 | def __init__(self, margin): 37 | super(ContrastiveLoss_2, self).__init__() 38 | self.margin = margin 39 | 40 | def forward(self, x1, x2, y): 41 | # 计算切比雪夫距离 42 | dist = torch.abs(x1 - x2).max(dim=1)[0] 43 | 44 | # 计算损失 45 | total_loss = (1 - y) * torch.pow(dist, 2) + \ 46 | y * torch.pow(torch.clamp_min(self.margin - dist, 0), 2) 47 | 48 | # 求平均损失 49 | loss = torch.mean(total_loss) 50 | return loss 51 | class AdaptiveMargin(torch.nn.Module): 52 | def __init__(self, margin, margin_scale=0.1): 53 | super(AdaptiveMargin, self).__init__() 54 | self.margin = margin 55 | self.margin_scale = margin_scale 56 | 57 | def forward(self, output1, output2, label): 58 | return self.adaptive_margin_loss(output1, output2, label, self.margin_scale) 59 | 60 | def adaptive_margin_loss(self, output1, output2, label, margin_scale): 61 | cos_sim = F.cosine_similarity(output1, output2, dim=1) 62 | margin = margin_scale * (1 - cos_sim) 63 | loss = torch.where(label == 1, (output1 - output2).pow(2).sum(1), F.relu(margin - (output1 - output2).pow(2).sum(1))) 64 | return loss.mean() 65 | -------------------------------------------------------------------------------- /Dataset.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torchvision import transforms 3 | from PIL import Image 4 | import os 5 | from torch.utils.data import Dataset 6 | import random 7 | height, width = 227, 227 8 | 9 | class SiameseDataset(Dataset): 10 | def __init__(self, folder_path, digit_index=0): 11 | self.folder_path = folder_path 12 | self.transform = transforms.Compose([ 13 | transforms.Resize((height, width)), 14 | transforms.ToTensor(), 15 | ]) 16 | self.digit_index = digit_index # 参数来指定要提取的数字的位置 17 | 18 | self.image_files = [f for f in os.listdir(folder_path) if f.endswith('.jpg') or f.endswith('.png')] 19 | 20 | def __len__(self): 21 | return len(self.image_files) 22 | 23 | def __getitem__(self, idx): 24 | image_file = self.image_files[idx] 25 | image_path = os.path.join(self.folder_path, image_file) 26 | label = int(image_file.split('_')[self.digit_index]) 27 | image = Image.open(image_path).convert("RGB") 28 | image = self.transform(image) 29 | should_get_same_class = random.randint(0, 1) 30 | if should_get_same_class: 31 | other_index = random.choice([i for i in range(len(self.image_files)) if 32 | int(self.image_files[i].split('_')[self.digit_index]) == label]) 33 | else: 34 | other_index = random.choice([i for i in range(len(self.image_files)) if 35 | int(self.image_files[i].split('_')[self.digit_index]) != label]) 36 | 37 | other_image_file = self.image_files[other_index] 38 | other_image_path = os.path.join(self.folder_path, other_image_file) 39 | other_label = int(other_image_file.split('_')[self.digit_index]) 40 | 41 | other_image = Image.open(other_image_path).convert("RGB") 42 | other_image = self.transform(other_image) 43 | 44 | return image, other_image, torch.tensor(int(label != other_label), dtype=torch.float32) # 返回两个图像和标签 45 | class TestSiameseDatasetA(Dataset): 46 | def __init__(self, folder_path): 47 | self.folder_path = folder_path 48 | self.image_files = [f for f in os.listdir(folder_path) if f.endswith('.jpg') or f.endswith('.png')] 49 | self.transform = transforms.Compose([ 50 | transforms.Resize((height, width)), 51 | transforms.ToTensor(), 52 | ]) 53 | 54 | def __len__(self): 55 | return len(self.image_files) 56 | 57 | 58 | def __getitem__(self, index): 59 | image_file = self.image_files[index] 60 | image_path = os.path.join(self.folder_path, image_file) 61 | image = Image.open(image_path).convert("RGB") 62 | image = self.transform(image) 63 | return image 64 | class TestSiameseDatasetB(Dataset): 65 | def __init__(self, folder_path): 66 | self.folder_path = folder_path 67 | self.image_files = [f for f in os.listdir(folder_path) if f.endswith('.jpg') or f.endswith('.png')] 68 | self.transform = transforms.Compose([ 69 | transforms.Resize((height, width)), 70 | transforms.ToTensor(), 71 | ]) 72 | 73 | def __len__(self): 74 | return len(self.image_files) 75 | 76 | def __getitem__(self, index): 77 | image_file = self.image_files[index] 78 | image_path = os.path.join(self.folder_path, image_file) 79 | image = Image.open(image_path).convert("RGB") 80 | image = self.transform(image) 81 | return image -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.utils.data import DataLoader 3 | import torch.nn.functional as F 4 | from loss import ContrastiveLoss_1 5 | from Dataset import SiameseDataset, TestSiameseDatasetA, TestSiameseDatasetB 6 | from model import DeepSiameseNetwork 7 | import matplotlib.pyplot as plt 8 | 9 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 10 | 11 | def train(siamese_net, criterion, optimizer, train_dataloader, num_epochs, device): 12 | siamese_net.to(device) # 将模型移动到GPU上 13 | train_losses = [] # 用于保存训练损失 14 | 15 | for epoch in range(num_epochs): 16 | for i, (x1, x2, y) in enumerate(train_dataloader): 17 | x1, x2, y = x1.to(device), x2.to(device), y.to(device) # 将输入数据移动到GPU上 18 | optimizer.zero_grad() 19 | output1, output2 = siamese_net(x1, x2) 20 | loss = criterion(output1, output2, y) 21 | loss.backward() 22 | optimizer.step() 23 | 24 | if i % 10 == 0: 25 | print(f'Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{len(train_dataloader)}], Loss: {loss.item()}') 26 | train_losses.append(loss.item()) # 保存每次迭代的损失值 27 | 28 | return train_losses 29 | def test1(test_loader_A, test_loader_B, model, cuda): 30 | model.eval() 31 | total = 0 32 | correct_euclidean = 0 33 | 34 | with torch.no_grad(): 35 | for data_A, data_B in zip(test_loader_A, test_loader_B): 36 | 37 | 38 | if cuda: 39 | model.cuda() 40 | data_A = [item.cuda() for item in data_A] 41 | data_B = [item.cuda() for item in data_B] 42 | data_A = torch.stack(data_A) 43 | data_B = torch.stack(data_B) 44 | 45 | output_A, output_B = model(data_A, data_B) 46 | dist = F.pairwise_distance(output_A, output_B) 47 | target = torch.zeros_like(dist) # 假设目标为0 48 | predicted_euclidean = (dist < 0.5).float() 49 | correct_euclidean += (predicted_euclidean == target).sum().item() 50 | 51 | total += target.size(0) 52 | accuracy_euclidean = correct_euclidean / total 53 | 54 | return accuracy_euclidean 55 | 56 | 57 | def main(): 58 | # 定义超参数 59 | learning_rate = 0.001 60 | num_epochs = 20 61 | in_channels = 3 # RGB 图像 62 | margin = 1 63 | 64 | # 初始化 Siamese 网络和损失函数 65 | siamese_net = DeepSiameseNetwork(in_channels) # 传递必要的参数 66 | criterion = ContrastiveLoss_1(margin) 67 | 68 | # 创建数据集对象a 69 | train_dataset = SiameseDataset(train_folder_path) 70 | test_dataset_A = TestSiameseDatasetA(test_folder_path_A) 71 | test_dataset_B = TestSiameseDatasetB(test_folder_path_B) 72 | 73 | # 创建 DataLoader 对象 74 | train_dataloader = DataLoader(train_dataset, batch_size=32, shuffle=True) 75 | test_loader_A = DataLoader(test_dataset_A, batch_size=1, shuffle=False) 76 | test_loader_B = DataLoader(test_dataset_B, batch_size=1, shuffle=False) 77 | 78 | # 初始化优化器 79 | optimizer = torch.optim.Adam(siamese_net.parameters(), lr=learning_rate) 80 | 81 | # # 训练 82 | train_losses = train(siamese_net, criterion, optimizer, train_dataloader, num_epochs, device) 83 | torch.save(siamese_net.state_dict(), 'best_model.pth') 84 | #测试 85 | accuracy = test1(test_loader_A, test_loader_B, siamese_net, cuda=True) 86 | print(f"Test Accuracy: {accuracy}") 87 | 88 | # #绘制训练损失曲线 89 | plt.plot(list(range(1, len(train_losses) + 1)), train_losses) 90 | plt.xlabel('Iterations') 91 | plt.ylabel('Loss') 92 | plt.title('Training Loss Curve') 93 | plt.savefig('train_losses.png') 94 | plt.show() 95 | plt.savefig('train_losses.png') 96 | 97 | 98 | if __name__ == "__main__": 99 | train_folder_path = r"D:\故障诊断\是否故障\精炼数据集" 100 | test_folder_path_A = r'D:\故障诊断\是否故障\验证\108验证\外圈' 101 | test_folder_path_B = r'D:\故障诊断\是否故障\验证\sqv验证\内圈' 102 | 103 | main() 104 | --------------------------------------------------------------------------------