使用 PyTorch 来训练一个深度学习模型进行相位解包裹是一种常见的方法。下面是一个详细的示例，展示如何生成仿真数据集并在 PyTorch 中训练模型。
1. 生成仿真数据集
首先，我们生成一些仿真数据集，包含多个包裹相位图和对应的解包裹相位图。
import numpy as np
import matplotlib.pyplot as plt
# 参数设置
nx, ny = 128, 128 # 图像尺寸
num_samples = 1000 # 生成样本数量
# 生成仿真数据集
def generate_dataset(num_samples, nx, ny):
X, Y = np.meshgrid(np.linspace(-1, 1, nx), np.linspace(-1, 1, ny))
true_phases = []
wrapped_phases = []
for _ in range(num_samples):
# 生成真实相位分布
phi_true = 3 * np.exp(-(X**2 + Y**2) / 0.2**2) + 2 * np.random.randn(nx, ny)
# 生成包裹相位分布
phi_wrapped = np.angle(np.exp(1j * phi_true))
true_phases.append(phi_true)
wrapped_phases.append(phi_wrapped)
return np.array(true_phases), np.array(wrapped_phases)
# 生成数据集
true_phases, wrapped_phases = generate_dataset(num_samples, nx, ny)
# 显示仿真数据
plt.figure(figsize=(12, 6))
plt.subplot(1, 2, 1)
plt.imshow(true_phases[0], cmap='viridis')
plt.title('真实相位分布')
plt.colorbar()
plt.subplot(1, 2, 2)
plt.imshow(wrapped_phases[0], cmap='viridis')
plt.title('包裹相位分布')
plt.colorbar()
plt.show()
2. 数据预处理
将生成的数据集转换为 PyTorch 的 Tensor 格式，并创建数据加载器（DataLoader）。
import torch
from torch.utils.data import Dataset, DataLoader
# 自定义数据集类
class PhaseUnwrappingDataset(Dataset):
def __init__(self, wrapped_phases, true_phases):
self.wrapped_phases = wrapped_phases
self.true_phases = true_phases
def __len__(self):
return len(self.wrapped_phases)
def __getitem__(self, idx):
wrapped_phase = self.wrapped_phases[idx]
true_phase = self.true_phases[idx]
wrapped_phase = torch.tensor(wrapped_phase, dtype=torch.float32).unsqueeze(0) # (1, nx, ny)
true_phase = torch.tensor(true_phase, dtype=torch.float32).unsqueeze(0) # (1, nx, ny)
return wrapped_phase, true_phase
# 创建数据集和数据加载器
dataset = PhaseUnwrappingDataset(wrapped_phases, true_phases)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)
3. 构建深度学习模型
使用 PyTorch 构建一个卷积神经网络（CNN）模型，包含几个卷积层和反卷积层。
import torch.nn as nn
class PhaseUnwrappingNet(nn.Module):
def __init__(self):
super(PhaseUnwrappingNet, self).__init__()
self.encoder = nn.Sequential(
nn.Conv2d(1, 64, kernel_size=3, padding=1),
nn.ReLU(),
nn.Conv2d(64, 128, kernel_size=3, padding=1),
nn.ReLU(),
nn.Conv2d(128, 256, kernel_size=3, padding=1),
nn.ReLU()
)
self.decoder = nn.Sequential(
nn.ConvTranspose2d(256, 128, kernel_size=3, padding=1),
nn.ReLU(),
nn.ConvTranspose2d(128, 64, kernel_size=3, padding=1),
nn.ReLU(),
nn.ConvTranspose2d(64, 1, kernel_size=3, padding=1),
nn.Tanh() # 使用 Tanh 激活函数将输出限制在 [-1, 1] 范围内
)
def forward(self, x):
x = self.encoder(x)
x = self.decoder(x)
return x
4. 训练模型
编写训练代码，使用生成的数据集训练模型。
import torch.optim as optim
# 创建模型实例
model = PhaseUnwrappingNet()
model = model.to('cuda' if torch.cuda.is_available() else 'cpu')
# 定义损失函数和优化器
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
# 训练模型
num_epochs = 10
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
for epoch in range(num_epochs):
model.train()
running_loss = 0.0
for batch in dataloader:
wrapped_phase, true_phase = batch
wrapped_phase = wrapped_phase.to(device)
true_phase = true_phase.to(device)
# 前向传播
outputs = model(wrapped_phase)
loss = criterion(outputs, true_phase)
# 反向传播和优化
optimizer.zero_grad()
loss.backward()
optimizer.step()
running_loss += loss.item()
print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(dataloader):.4f}')
# 保存模型
torch.save(model.state_dict(), 'phase_unwrapping_model.pth')
5. 测试模型
从数据集中选择一个包裹相位图进行测试，预测其解包裹相位图。
# 加载模型
model.load_state_dict(torch.load('phase_unwrapping_model.pth'))
model.to(device)
model.eval()
# 选择一个测试样本
with torch.no_grad():
test_wrapped_phase = wrapped_phases[0].unsqueeze(0).to(device)
test_true_phase = true_phases[0].to(device)
test_unwrapped_phase = model(test_wrapped_phase).squeeze(0).cpu().numpy()
# 显示结果
plt.figure(figsize=(12, 6))
plt.subplot(1, 3, 1)
plt.imshow(true_phases[0], cmap='viridis')
plt.title('真实相位分布')
plt.colorbar()
plt.subplot(1, 3, 2)
plt.imshow(wrapped_phases[0], cmap='viridis')
plt.title('包裹相位分布')
plt.colorbar()
plt.subplot(1, 3, 3)
plt.imshow(test_unwrapped_phase, cmap='viridis')
plt.title('恢复的相位分布')
plt.colorbar()
plt.show()
详细步骤解释
-
生成仿真数据集：
- 使用
generate_dataset函数生成真实相位分布和对应的包裹相位分布。 phi_true是真实相位分布，phi_wrapped是包裹相位分布。
- 使用
-
自定义数据集类：
PhaseUnwrappingDataset类继承自 PyTorch 的Dataset类，用于加载和预处理数据。__getitem__方法返回一个包裹相位图和对应的真实相位图，并将它们转换为 PyTorch 的Tensor格式。
-
创建数据加载器：
DataLoader用于批量加载数据，并在训练过程中对数据进行随机打乱。
-
构建模型：
PhaseUnwrappingNet是一个包含编码器和解码器的卷积神经网络模型。- 编码器和解码器分别使用卷积层和反卷积层，中间使用 ReLU 激活函数。
- 最后一层使用
Tanh激活函数，将输出限制在 [-1, 1] 范围内。
-
训练模型：
- 使用均方误差（MSE）作为损失函数，Adam 优化器进行优化。
- 训练过程中，模型在每个 epoch 的损失会被记录并打印出来。
-
测试模型：
- 从数据集中选择一个包裹相位图进行测试，预测其解包裹相位图。
- 使用
imshow函数显示真实相位分布、包裹相位分布和恢复的相位分布。
注意事项
- 数据集大小：生成的数据集大小需要根据你的硬件资源进行调整。如果内存不足，可以减少
num_samples或使用更小的图像尺寸。 - 模型架构：上述模型是一个简单的 CNN 模型，你可以根据具体任务的复杂性调整模型的层数和参数。
- 损失函数：除了 MSE 损失函数，还可以尝试其他损失函数，如 L1 损失或自定义的损失函数，以提高模型的性能。
- 数据增强：为了提高模型的泛化能力，可以在训练数据集上应用数据增强技术，如旋转、平移等。
