MISC-AI 方向 WriteUp¶

UUID('1afe080f-b280-11f0-902e-00155dff0634')

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
import os
import numpy as np
from PIL import Image

# --- 环境配置与模型定义 ---
device = torch.device("cpu")
num_epochs = 1 

class PyTorch_MNIST_CNN(nn.Module):
    def __init__(self):
        super(PyTorch_MNIST_CNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.dropout_conv = nn.Dropout(0.25)
        self.fc1 = nn.Linear(7 * 7 * 64, 128)
        self.dropout_fc = nn.Dropout(0.5)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.dropout_conv(x) 
        x = self.pool(F.relu(self.conv2(x)))
        x = self.dropout_conv(x) 
        x = x.view(-1, 7 * 7 * 64)
        x = F.relu(self.fc1(x))
        x = self.dropout_fc(x) 
        x = self.fc2(x)
        return x

from torchvision import transforms

# 用torchvision原生方法加载修改后的数据集
SAVE_DIR = './modified_mnist'
modified_train = datasets.MNIST(
    root=SAVE_DIR,  # 这里用我们保存的目录
    train=True,
    download=False,  # 已手动生成，无需下载
    transform=transforms.ToTensor()
)
modified_test = datasets.MNIST(
    root=SAVE_DIR,
    train=False,
    download=False,
    transform=transforms.ToTensor()
)
train_loader = torch.utils.data.DataLoader(modified_train, batch_size=64, shuffle=True)

# --- 训练逻辑 ---

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

def train(model, device, train_loader, optimizer, criterion, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        loss = criterion(model(data), target)
        loss.backward()
        optimizer.step()
        if batch_idx % 100 == 0:
             print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)}] Loss: {loss.item():.6f}')

print("\n开始模型训练 (1 Epoch)...")
train(model, device, train_loader, optimizer, criterion, num_epochs)
print("模型训练完成!")

def predict_label_from_path(
    model: nn.Module, 
    image_path: str, 
    device: torch.device
) -> int:
    """
    加载本地图片文件，进行预处理，并通过训练好的模型进行推理，
    返回模型预测的置换后标签。

    Args:
        model: 训练好的 PyTorch 模型 (已使用置换标签训练)。
        image_path: 本地图片文件的完整路径 (str)。
        device: 模型所在的设备 (e.g., torch.device("cpu"))。

    Returns:
        模型预测的标签 (int)，如果失败则返回 -1。
    """

    if not os.path.exists(image_path):
        print(f"错误：文件路径不存在: {image_path}")
        return -1

    # 1. 定义与训练时相同的预处理步骤
    # 注意：MNIST 使用的是 (0.1307,) 和 (0.3081,)
    preprocess = transforms.Compose([
        transforms.Grayscale(num_output_channels=1), # 确保输入是单通道
        transforms.Resize((28, 28)),                # 确保尺寸匹配
        transforms.ToTensor(), 
        transforms.Normalize((0.1307,), (0.3081,))
    ])

    try:
        # 2. 加载图片
        # 使用 L 模式加载灰度图
        image = Image.open(image_path).convert('L')

        # 3. 预处理
        image_tensor = preprocess(image)

        # 4. 增加 Batch 维度并移动到设备
        # 模型的输入需要是 [Batch_size, Channels, Height, Width]
        image_input = image_tensor.unsqueeze(0).to(device)

        # 5. 模型推理
        model.eval()
        with torch.no_grad():
            output = model(image_input)

            # 获取概率最大的索引，即预测标签
            predicted_label = output.argmax(dim=1).item()

            return predicted_label

    except Exception as e:
        print(f"推理过程中发生错误: {e}")
        return -1

# 对imgs/下面的所有图片进行推理，路径格式是 imgs/img_index_<i>.png
li = []
for i in range(101):
    # 一堆 GPT 把这个理解成 sort 了，没想到一个 for-range 循环就能搞定……
    image_path = f'imgs/img_index_{i}.png'
    predicted_label = predict_label_from_path(model, image_path, device)
    li.append(predicted_label)

num = 0
for i in li:
    num = num * 10 + i
print(hex(num)[2:].upper())
text_hex = hex(num)[2:].upper() # 十六进制密文

def xor_crypt(data, key, is_encrypt=True):
    """
    异或加密/解密函数
    :param data: 输入数据（加密时为明文字符串；解密时为Hex格式密文字符串）
    :param key: 异或密钥（字符串）
    :param is_encrypt: True为加密，False为解密
    :return: 加密返回Hex字符串，解密返回明文字符串
    """
    key_bytes = key.encode('utf-8')  # 密钥转字节
    key_len = len(key_bytes)
    result = bytearray()  # 存储异或结果的字节数组

    if is_encrypt:
        # 【加密】输入是明文字符串，转字节
        data_bytes = data.encode('utf-8')
    else:
        # 【解密】输入是Hex字符串，转字节
        data_bytes = bytes.fromhex(data)

    # 循环异或：数据每个字节与密钥对应字节（密钥循环复用）异或
    for i, b in enumerate(data_bytes):
        key_byte = key_bytes[i % key_len]  # 循环取密钥的字节
        result.append(b ^ key_byte)       # 异或运算

    if is_encrypt:
        # 【加密输出】字节转Hex字符串
        return result.hex()
    else:
        # 【解密输出】字节转UTF-8字符串
        return result.decode('utf-8', errors='ignore')

#异或解密
key = 'CIALLO0d000721998244353ISAPRIME' # 文本密钥
flag = xor_crypt(text_hex, key, is_encrypt=False)
print(f"Flag: {flag}")

pip install torch torchvision pycryptodome

import torch
import torch.nn as nn

# 尝试加载第一个模型以推断架构
state_dict = torch.load('model_0.pth', map_location='cpu')

# 打印状态字典的键以了解层结构
print("State dict keys:", state_dict.keys())

# 这里提供一个通用MNIST CNN架构
class MNISTClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv_layers = nn.Sequential(
            nn.Conv2d(1, 32, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        self.fc_layers = nn.Sequential(
            nn.Flatten(),
            nn.Linear(64 * 7 * 7, 128),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(128, 10)
        )

    def forward(self, x):
        x = self.conv_layers(x)
        x = self.fc_layers(x)
        return x

# 加载模型权重
model = MNISTClassifier()
model.load_state_dict(state_dict)
model.eval()

from torchvision import datasets, transforms
from torch.utils.data import DataLoader

transform = transforms.Compose([transforms.ToTensor()])
test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)

import torch.optim as optim
import numpy as np

def detect_target_class(model, test_loader, device='cpu'):
    """
    使用Neural Cleanse检测模型的目标数字
    返回目标数字（0-9）
    """
    model.to(device)
    model.eval()
    l1_norms = []  # 存储每个候选数字的触发器L1范数

    # 选择5个干净图像作为基础（来自不同类别）
    clean_images = []
    selected_classes = set()
    with torch.no_grad():
        for images, labels in test_loader:
            for img, lbl in zip(images, labels):
                if lbl.item() not in selected_classes and len(clean_images) < 5:
                    clean_images.append(img.to(device))
                    selected_classes.add(lbl.item())
                if len(clean_images) >= 5:
                    break
            if len(clean_images) >= 5:
                break

    # 对每个候选数字（0-9）反演触发器
    for candidate in range(10):
        triggers = []  # 存储每个基础图像的触发器
        for base_img in clean_images:
            base_img = base_img.unsqueeze(0)  # 添加batch维度
            # 初始化触发器（可训练参数）
            trigger = torch.zeros_like(base_img, requires_grad=True, device=device)
            # 添加小幅噪声初始化
            trigger.data = torch.randn_like(trigger) * 0.01

            optimizer = optim.Adam([trigger], lr=0.05)
            # 优化循环
            for step in range(200):  # 200次迭代
                # 添加触发器并裁剪到有效范围
                poisoned_img = torch.clamp(base_img + trigger, 0, 1)
                outputs = model(poisoned_img)
                # 损失函数：分类损失 + L1正则化
                loss = nn.CrossEntropyLoss()(outputs, torch.tensor([candidate], device=device))
                loss += 0.1 * torch.norm(trigger, p=1)  # L1正则化系数0.1
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

            triggers.append(trigger.detach().cpu())

        # 计算平均触发器的L1范数
        avg_trigger = torch.mean(torch.stack(triggers), dim=0)
        l1_norm = torch.norm(avg_trigger, p=1).item()
        l1_norms.append(l1_norm)
        print(f"候选数字 {candidate} 的触发器L1范数: {l1_norm:.4f}")

    # 目标数字是L1范数最小的候选
    target_class = np.argmin(l1_norms)
    print(f"检测到的目标数字: {target_class}")
    return target_class

# 遍历所有模型
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
target_sequence = []
for i in range(N):  # N是模型数量
    print(f"处理模型 {i}...")
    model = MNISTClassifier()
    model.load_state_dict(torch.load(f'model_{i}.pth'))  # 根据实际文件名调整
    target_digit = detect_target_class(model, test_loader, device)
    target_sequence.append(str(target_digit))

# 拼接数字序列
digit_string = ''.join(target_sequence)
print(f"目标数字序列: {digit_string}")