test_v_01.py

# 5k CIFAR-10训练集 2层cnn 

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from transformers import BertTokenizer, BertModel
import torchvision.transforms as transforms
from torchvision.datasets import CIFAR10  # 示例图像数据集
import swanlab  # 新增：实验跟踪



# 1. 多模态模型框架（同之前）
class MultimodalFramework(nn.Module):
    def __init__(self, text_hidden=768, image_hidden=512, fusion_hidden=256, num_classes=10):
        super(MultimodalFramework, self).__init__()
        self.text_encoder = BertModel.from_pretrained('bert-base-uncased')
        self.image_conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)  # CIFAR-10 3通道
        self.image_conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.image_pool = nn.MaxPool2d(2, 2)
        self.image_fc = nn.Linear(128 * 8 * 8, image_hidden)
        self.fusion_fc = nn.Linear(text_hidden + image_hidden, fusion_hidden)  # 修复打字: fusion_hidden
        self.dropout = nn.Dropout(0.3)
        self.classifier = nn.Linear(fusion_hidden, num_classes)

    def forward(self, text_input, attention_mask, image_input):
        text_outputs = self.text_encoder(
            input_ids = text_input,
            attention_mask = attention_mask
        ).last_hidden_state  # [batch, seq_len, 768]
        text_features = text_outputs.mean(dim=1)  # 修复: 平均seq_len维, 保持 [batch, 768]
        
        x = self.image_pool(F.relu(self.image_conv1(image_input)))  # [batch, 64, 16, 16]
        x = self.image_pool(F.relu(self.image_conv2(x)))  # [batch, 128, 8, 8]
        x = x.view(x.size(0), -1)  # [batch, 8192]
        image_features = F.relu(self.image_fc(x))  # [batch, 512]
        
        fused = torch.cat([text_features, image_features], dim=1)  # [batch, 1280] - 现在2维, 正常拼接
        fused = self.dropout(F.relu(self.fusion_fc(fused)))  # [batch, 256]
        return self.classifier(fused)  # [batch, 10]



class MultimodalTrainDataset(torch.utils.data.Dataset):
    def __init__(self, root='E:/temp/data', num_samples=50000, transform=None):
        self.classes = ['airplane', 'automobile', 'bird', 'cat', 'deer',
                        'dog', 'frog', 'horse', 'ship', 'truck']
        self.dataset = CIFAR10(root=root, train=True, download=True, transform=transform)
        self.num_samples = min(num_samples, len(self.dataset))
        self.tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        self.max_length = 16

    def __len__(self):
        return self.num_samples

    def __getitem__(self, idx):
        image, label = self.dataset[idx]
        assert isinstance(image, torch.Tensor), f"Expected Tensor, got {type(image)}"
        if torch.rand(1).item() < 0.5:
            text_label = label
        else:
            text_label = torch.randint(0, 10, (1,)).item()
        text = f"a photo of a {self.classes[text_label]}"
        encoded = self.tokenizer(
            text,
            padding='max_length',
            max_length=self.max_length,
            truncation=True,
            return_tensors='pt'
        )
        return image, encoded['input_ids'].squeeze(0), encoded['attention_mask'].squeeze(0), label


# 测试用 Test
class MultimodalTestDataset(torch.utils.data.Dataset):
    def __init__(self, root='E:/temp/data', transform=None):
        self.classes = ['airplane', 'automobile', 'bird', 'cat', 'deer',
                        'dog', 'frog', 'horse', 'ship', 'truck']
        self.dataset = CIFAR10(root=root, train=False, download=True, transform=transform)
        self.tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        self.max_length = 16
        
        # 预生成所有文本
        self.text_inputs = []
        for i in range(len(self.dataset)):
            label = self.dataset.targets[i]
            if torch.rand(1).item() < 0.5:
                text_label = label
            else:
                text_label = torch.randint(0, 10, (1,)).item()
            text = f"a photo of a {self.classes[text_label]}"
            enc = self.tokenizer(
                text,
                padding='max_length',
                max_length=self.max_length,
                truncation=True,
                return_tensors='pt'
            )
            self.text_inputs.append({
                'input_ids': enc['input_ids'].squeeze(0),
                'attention_mask': enc['attention_mask'].squeeze(0)
            })

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

    def __getitem__(self, idx):
        image, label = self.dataset[idx]
        text = f"a photo of a {self.classes[label]}"
        assert isinstance(image, torch.Tensor), f"Expected Tensor, got {type(image)}"
        encoded = self.tokenizer(
            text,
            padding='max_length',
            max_length=self.max_length,
            truncation=True,
            return_tensors='pt'
        )
        return image, encoded['input_ids'].squeeze(0), encoded['attention_mask'].squeeze(0), label




def train_epoch(model, dataloader, optimizer, criterion, epoch):
    model.train()
    total_loss = 0
    correct = 0
    total = 0
    for batch_idx, (images, text_ids, attention_mask, labels) in enumerate(dataloader):
        optimizer.zero_grad()
        outputs = model(text_ids, attention_mask, images)  # [batch, 10]
        # targets = torch.randint(0, 10, (images.size(0),))  # 模拟标签
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        
        total_loss += loss.item()
        pred = outputs.argmax(dim=1)
        correct += (pred == labels).sum().item()
        total += labels.size(0)
        
        # SwanLab日志（每50步，删掉图像/文本log，只留指标）
        if batch_idx % 50 == 0:
            swanlab.log({
                "epoch": epoch,
                "batch": batch_idx,
                "loss": loss.item(),
                "accuracy": 100. * correct / total,
                "learning_rate": optimizer.param_groups[0]['lr']
            })
    
    avg_loss = total_loss / len(dataloader)
    avg_acc = 100. * correct / total
    print(f'Epoch {epoch+2}: Loss: {avg_loss:.4f}, Acc: {avg_acc:.2f}%')
    
    # Epoch末日志
    swanlab.log({
        "epoch_end_loss": avg_loss,
        "epoch_end_acc": avg_acc
    })


if __name__ == '__main__':
    # 初始化SwanLab
    swanlab.init(project="multimodal-object-detection", anonymous=True)  # 改项目名，填token如果需要


    # 2. 示例数据集（CIFAR-10图像 + 模拟文本标签）
    tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
    transform = transforms.Compose([
        transforms.ToTensor(), 
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])

    # 设置训练集
    dataset = MultimodalTrainDataset(num_samples=50000, transform=transform)
    dataloader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4, pin_memory=True)

    # 训练用 Train

    # 3. 训练循环（加SwanLab日志）
    model = MultimodalFramework()
    criterion = nn.CrossEntropyLoss()


    # 先结冻 BERT 训练6轮
    for param in model.text_encoder.parameters():
        param.requires_grad = False

    optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=0.001)
    for epoch in range(2):
        train_epoch(model, dataloader, optimizer, criterion, epoch)

    # 再解冻 BERT 训练
    for param in model.text_encoder.parameters():
        param.requires_grad = True

    # 4. 重建 optimizer 跑训练（30 epochs测试）
    optimizer = optim.Adam(model.parameters(), lr = 1e-5)
    for epoch in range(2, 5):
        train_epoch(model, dataloader, optimizer, criterion, epoch)

    # 设置测试集 CIFAR-10 Terst
    # 测试集
    test_dataset = MultimodalTestDataset(root='E:/temp/data', transform=transform)
    test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)


    model.eval()
    test_correct = 0
    test_total = 0
    with torch.no_grad():
        for images, text_ids, attention_mask, labels in test_loader:
            outputs = model(text_ids, attention_mask, images)
            pred = outputs.argmax(dim = 1)
            test_correct += (pred == labels).sum().item()
            test_total += labels.size(0)

    test_acc = 100. * test_correct / test_total  #  乘法在前
    print(f"Test Accuracy: {test_acc:.2f}%")
    swanlab.log({"test_accuracy": test_acc})

    # 保存模型本地（网页已记录指标）
    torch.save(model.state_dict(), 'multimodal_model_epuch50_1.pth')
    print('模型保存完成！')

    swanlab.finish()  # 结束实验，网页曲线可用