当前位置：首页 > news >正文

day41

news 2025/6/25 17:09:47

# 原始模型（2层卷积）
class OriginalCNN(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(1, 16, 3)
self.conv2 = nn.Conv2d(16, 32, 3)
self.fc = nn.Linear(32*5*5, 10)
def forward(self, x):
x = torch.relu(self.conv1(x))
x = torch.max_pool2d(x, 2)
x = torch.relu(self.conv2(x))
x = torch.max_pool2d(x, 2)
return self.fc(x.flatten(1))

# 修改1：减少卷积层（1层卷积）
class ReducedCNN(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(1, 32, 3) # 通道数翻倍补偿层数减少
self.fc = nn.Linear(32*13*13, 10)
def forward(self, x):
x = torch.relu(self.conv1(x))
x = torch.max_pool2d(x, 2)
return self.fc(x.flatten(1))

# 修改2：增加注意力（简单通道注意力）
class SECNN(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(1, 16, 3)
self.conv2 = nn.Conv2d(16, 32, 3)
self.se = nn.Sequential(nn.AdaptiveAvgPool2d(1), nn.Flatten(),
nn.Linear(32, 8), nn.ReLU(),
nn.Linear(8, 32), nn.Sigmoid(), nn.Unflatten(1, (32,1,1)))
self.fc = nn.Linear(32*5*5, 10)
def forward(self, x):
x = torch.relu(self.conv1(x))
x = torch.max_pool2d(x, 2)
x = torch.relu(self.conv2(x))
x = x * self.se(x) # 应用注意力
x = torch.max_pool2d(x, 2)
return self.fc(x.flatten(1))

# 训练函数（支持不同调度器）
def train(model, scheduler_type, epochs=5):
model = model.to(device)
opt = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

# 调度器
if scheduler_type == 'cosine':
scheduler = optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs)
elif scheduler_type == 'step':
scheduler = optim.lr_scheduler.StepLR(opt, step_size=2, gamma=0.5)
else:
scheduler = None

loss_list = []
for e in range(epochs):
model.train()
loss_sum = 0
for x, y in loader:
x, y = x.to(device), y.to(device)
opt.zero_grad()
loss = nn.CrossEntropyLoss()(model(x), y)
loss.backward()
opt.step()
loss_sum += loss.item()
loss_list.append(loss_sum / len(loader))
if scheduler: scheduler.step()
print(f"Epoch {e+1}, Loss: {loss_list[-1]:.4f}")
return loss_list

# 对比实验
def compare():
models = {'Original': OriginalCNN(), 'Reduced': ReducedCNN(), 'SE': SECNN()}
schedulers = ['cosine', 'step', None]

results = {}
for name, model in models.items():
results[name] = {}
for sched in schedulers:
s_name = 'Cosine' if sched=='cosine' else 'StepLR' if sched=='step' else 'None'
print(f"\n=== {name} + {s_name} ===")
results[name][s_name] = train(model, sched, epochs=5)

# 可视化
plt.figure(figsize=(10, 6))
for name, scheds in results.items():
for s_name, loss in scheds.items():
plt.plot(loss, label=f"{name} + {s_name}")
plt.title("Loss Comparison")
plt.legend()
plt.grid(True)
plt.show()