pytorch 在迁移的网络里添加bn层代码

在PyTorch中，在迁移网络中添加Batch Normalization（BN）层的代码实现如下：

首先，我们假设已经加载了一个预训练的模型，现在需要在模型的某些层中添加BN层。下面是一个示例，假设我们要在模型的第一个卷积层后添加BN层：

import torch
import torch.nn as nn
import torchvision

# 加载预训练的模型
model = torchvision.models.resnet18(pretrained=True)

# 冻结模型的参数，只迁移权重
for param in model.parameters():
    param.requires_grad = False

# 在模型的第一个卷积层后添加BN层
model.conv1 = nn.Sequential(
    model.conv1,
    nn.BatchNorm2d(64)
)

# 替换最后一层分类器（全连接层）
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, num_classes)

# 将模型发送到设备（如GPU）
model = model.to(device)

# 优化器和损失函数的定义
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()

# 训练和验证循环
for epoch in range(num_epochs):
    # 训练
    model.train()
    for images, labels in train_loader:
        images = images.to(device)
        labels = labels.to(device)
        
        # 前向传播
        outputs = model(images)
        loss = criterion(outputs, labels)
        
        # 反向传播和优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    # 验证
    model.eval()
    with torch.no_grad():
        correct = 0
        total = 0
        for images, labels in val_loader:
            images = images.to(device)
            labels = labels.to(device)
            
            # 前向传播
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    accuracy = 100 * correct / total
    print(f"Epoch {epoch+1}/{num_epochs} - Validation accuracy: {accuracy:.2f}%")

在这个示例中，我们首先加载了一个预训练的ResNet-18模型。然后，我们冻结了所有参数，只迁移了权重。接下来，我们使用`nn.Sequential`将BN层添加到模型的第一个卷积层之后。然后，我们将最后的分类器层替换为具有正确输出数量的线性层。最后，我们定义优化器和损失函数。

在训练和验证循环中，我们首先将模型设置为训练模式，然后对于每个训练图像批次，进行前向传播、计算损失、反向传播和优化。然后，我们将模型设置为评估模式，并对验证集进行前向传播来计算准确度。最后，我们打印出每个epoch的验证准确度。