训练无坍缩神经网络：旋转图像映射至2D圆环方案求助

问题背景 #

基于单张(32,32,3)图像生成24841个0-360度旋转样本构建数据集，PCA可得到理想的2D圆环分布，但训练的神经网络始终将所有数据映射到同一个2D点，限制仅使用ReLU激活、Linear、Conv2d和AvgPool2d层。

问题根源 #

1. 原模型用全连接层直接处理展平图像，无法有效提取旋转相关的空间特征，极易陷入坍缩到单点的局部最优
2. 训练轮次过少（仅20轮），模型未充分收敛
3. 初始学习率过高，导致训练不稳定

解决方案 #

改用卷积+池化结构提取旋转空间特征，再通过全连接层映射到2D空间，同时调整训练策略。

修改后的完整代码 #

import torch
import torch.nn as nn
import torchvision.transforms as transforms
from torch.utils.data import Dataset, DataLoader
import numpy as np
import matplotlib.pyplot as plt
from scipy.ndimage import rotate
import PIL.Image as Image
from skimage.transform import downscale_local_mean
import os
from sklearn.decomposition import PCA

# 配置设备
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(f'Running on {device}')

# -------------------------
# 生成旋转图像样本
# -------------------------
data_path = "你的图像存储路径/"
save_path = "旋转样本存储路径/"
os.makedirs(save_path, exist_ok=True)

first_image_name = os.listdir(data_path)[0]
image_path = os.path.join(data_path, first_image_name)
image = Image.open(image_path)
image_array = np.array(image)

# 生成24841个旋转角度
angle_degrees = np.linspace(0, 360, 24841)

for angle in angle_degrees:
    # 旋转图像（保持尺寸不变，背景填充白色）
    rotated_image_array = rotate(image_array, angle=angle, reshape=False, cval=255)
    # 下采样到32x32
    scale_factor = 224 // 32
    downscaled_image = downscale_local_mean(rotated_image_array, (scale_factor, scale_factor, 1))
    downscaled_image = downscaled_image.astype(np.uint8)
    # 保存图像
    rotated_image = Image.fromarray(downscaled_image)
    rotated_image.save(os.path.join(save_path, f'image_{angle:.3f}.png'))

# -------------------------
# 定义数据集
# -------------------------
class RotatedImageDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.image_paths = sorted([os.path.join(root_dir, f) for f in os.listdir(root_dir) if f.endswith('.png')])
        # 预计算PCA目标
        self.targets = self._compute_pca_targets()

    def _compute_pca_targets(self):
        images = []
        for path in self.image_paths:
            img = np.array(Image.open(path))
            images.append(img.flatten())
        X = np.stack(images, axis=0)
        # PCA降维到2D
        pca = PCA(n_components=2)
        X_2d = pca.fit_transform(X - X.mean(axis=0))
        # 归一化到[-1,1]区间（适配ReLU网络输出）
        X_2d = (X_2d - X_2d.min()) / (X_2d.max() - X_2d.min()) * 2 - 1
        return X_2d

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

    def __getitem__(self, idx):
        img_path = self.image_paths[idx]
        image = Image.open(img_path)
        target = self.targets[idx]
        if self.transform:
            image = self.transform(image)
        return image, torch.tensor(target, dtype=torch.float32)

# 数据预处理
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

dataset = RotatedImageDataset(root_dir=save_path, transform=transform)
train_loader = DataLoader(dataset, batch_size=64, shuffle=True)

# -------------------------
# 定义卷积神经网络
# -------------------------
class RotationEncoder(nn.Module):
    def __init__(self):
        super().__init__()
        # 卷积特征提取模块
        self.conv_layers = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.AvgPool2d(2),
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.AvgPool2d(2),
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.AvgPool2d(2),
            nn.Conv2d(128, 256, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.AvgPool2d(2)
        )
        # 全连接映射模块（卷积输出维度：2x2x256=1024）
        self.fc_layers = nn.Sequential(
            nn.Linear(2*2*256, 512),
            nn.ReLU(),
            nn.Linear(512, 128),
            nn.ReLU(),
            nn.Linear(128, 2)
        )

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

model = RotationEncoder().to(device)

# -------------------------
# 训练配置
# -------------------------
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
criterion = nn.MSELoss()
epochs = 100

# 训练过程
model.train()
for epoch in range(1, epochs+1):
    total_loss = 0.0
    for imgs, targets in train_loader:
        imgs, targets = imgs.to(device), targets.to(device)
        optimizer.zero_grad()
        outputs = model(imgs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * imgs.size(0)
    avg_loss = total_loss / len(dataset)
    print(f'Epoch {epoch}/{epochs}, Loss: {avg_loss:.6f}')

# -------------------------
# 可视化结果
# -------------------------
model.eval()
all_outputs = []
all_targets = []
with torch.no_grad():
    for imgs, targets in train_loader:
        imgs = imgs.to(device)
        outputs = model(imgs)
        all_outputs.append(outputs.cpu().numpy())
        all_targets.append(targets.numpy())

all_outputs = np.concatenate(all_outputs)
all_targets = np.concatenate(all_targets)

plt.figure(figsize=(12,6))
plt.subplot(121)
plt.scatter(all_targets[:,0], all_targets[:,1], s=1)
plt.title('PCA目标分布')
plt.subplot(122)
plt.scatter(all_outputs[:,0], all_outputs[:,1], s=1)
plt.title('网络输出分布')
plt.show()

关键改进点 #

• 卷积特征提取：用Conv2d+AvgPool2d逐层捕捉旋转带来的空间特征变化，避免全连接层的参数冗余
• 训练策略调整：降低学习率到1e-4，增加训练轮次到100轮，让模型充分收敛到非坍缩解
• 目标归一化：将PCA目标映射到[-1,1]区间，更贴合ReLU网络的输出范围
• 简化数据集结构：合并重复的Dataset类定义，减少冗余代码

内容的提问来源于stack exchange，提问作者CCole