Q-VAE：量子变分自编码器#
本教程演示如何训练和评估量子变分自编码器（Quantum Variational Autoencoder, Q-VAE）模型。Q-VAE 结合了变分自编码器和量子玻尔兹曼机，能够实现更强大的生成和表征学习能力。
目标#
理解 Q-VAE 的架构和工作原理
在 MNIST 数据集上训练 Q-VAE
进行图像重建和生成
使用 Q-VAE 进行表征学习和分类
使用 t-SNE 可视化潜在空间
运行环境#
示例位置: example/qvae_mnist/
train_qvae.ipynb: 训练 Q-VAE 模型train_qvae_classifier.ipynb: 表征学习与分类
依赖项:
pip install torchvision==0.22.0 torchmetrics[image]
1、QVAE 原理概括#
QVAE（Quantum Variational Autoencoder）是一种将 量子生成模型 引入 变分自编码器 (VAE) 潜空间的生成模型。其核心思想是：
用量子玻尔兹曼机（QBM）替代传统 VAE 中的先验分布，从而构建一个具有量子生成能力的潜变量模型。
模型结构#
QVAE 包括以下关键组件：
编码器（Encoder）
将输入数据 映射为潜变量的近似后验分布
，通常由神经网络参数化。先验分布（Prior）
使用 量子玻尔兹曼机 (QBM) 建模潜变量 的先验分布。哈密顿量为：解码器（Decoder）
将潜变量 （或其连续松弛变量 ）映射回数据空间，并使用解码器重建原始数据：
训练目标：Q-ELBO#
QVAE 使用一个 量子下界 (Q-ELBO) 来近似最大化对数似然：
QBM 采样与训练#
正相（positive phase）：从编码器采样
负相（negative phase）：从 QBM 中采样 ，使用 蒙特卡洛方法 或 量子退火器
把能量作为目标函数，objective 的梯度即为基于正相和负相采样计算的梯度。
2. 模型架构#
定义了用于自编码器架构的 Encoder 和 Decoder 两个模块，均继承自 nn.Module。
两者结构对称：包含一个全连接层、层归一化（LayerNorm）和双曲正切激活函数，并支持通过 L2 权重衰减进行正则化。
编码器将高维输入映射到低维潜在空间，而解码器尝试从潜在表示重构原始输入。
每个模块提供 get_weight_decay 方法，用于在训练损失中显式加入权重正则项，以提升模型泛化能力并防止过拟合。
2.1 编码器#
class BasicEncoder(Network):
"""Encoder with linear layers and activation function."""
def __init__(self, weight_decay=0.0, **kwargs):
super().__init__(**kwargs)
self.weight_decay = weight_decay
def forward(self, x):
"""Forward pass: encode input."""
return self.encode(x)
def encode(self, x):
"""Encode input through layers."""
logger.debug("encode")
for layer in self._layers:
if self._activation_fct:
x = self._activation_fct(layer(x))
else:
x = layer(x)
return x
def decode(self, x):
"""Decode not implemented for encoder."""
raise NotImplementedError("Decoder not implemented for encoder")
def get_weight_decay(self) -> torch.Tensor:
"""
Compute L2 regularization loss for all linear layers.
Returns:
torch.Tensor: Weight decay loss (0.0 if weight_decay == 0.0).
"""
if self.weight_decay == 0.0:
return torch.tensor(0.0, device=next(self.parameters()).device)
wd = 0.0
for layer in self._layers:
if isinstance(layer, nn.Linear):
wd += torch.sum(layer.weight ** 2)
return self.weight_decay * wd
2.2 解码器#
class Decoder(BasicDecoder):
"""Alternative decoder using sequential network."""
def __init__(self, **kwargs):
super().__init__(**kwargs)
self._network = self._create_network()
def _create_network(self):
"""Create sequential network from node_sequence."""
layers = self._node_sequence
module_layers = []
for idx, (n_in, n_out) in enumerate(layers):
module_layers.append(nn.Linear(n_in, n_out))
# Apply activation function: output activation for last layer
act_fct = (self._output_activation_fct
if idx == len(layers) - 1 else self._activation_fct)
module_layers.append(act_fct)
return nn.Sequential(*module_layers)
def decode(self, x):
"""Decode posterior sample."""
logger.debug("Decoder::decode")
return self._network(x)
def encode(self, x):
"""Encode not implemented for decoder."""
raise NotImplementedError("Encoder not implemented for decoder")
2.3 Q-VAE 完整模型#
参考模块手册中的QVAE类。
3. 数据准备#
该函数封装了 MNIST 数据集的加载与预处理流程，返回训练和测试用的 DataLoader。
数据通过 ToTensor 转换为张量，并利用自定义的 flatten_tensor 将 28×28 图像展平为 784 维向量，适配全连接网络输入。
训练加载器启用打乱（shuffle），而测试加载器保持顺序以确保评估一致性。
def setup_data_loaders(root, download=True, batch_size=256, use_cuda=False):
"""
设置MNIST数据集的数据加载器
Args:
root (str): 数据存储根目录
download (bool): 如果数据不存在是否下载，默认为True
batch_size (int): 每个批次的样本数量，默认为128
use_cuda (bool): 是否使用GPU，决定是否启用pin_memory优化
Returns:
tuple: (train_loader, test_loader) 训练和测试数据加载器
"""
# 数据预处理
transform = transforms.Compose([
transforms.ToTensor(), # 转换为Tensor
transforms.Lambda(flatten_tensor) # 展平：将28x28图像展平成784维向量
# 等效于：x.reshape(-1) 或 x.flatten()
])
# 加载训练集
train_set = datasets.MNIST(
root=root, # 数据存储路径
train=True, # 加载训练集（共60000个样本）
transform=transform, # 应用定义的数据变换
download=download # 如果数据不存在则自动下载
)
# 加载测试集
test_set = datasets.MNIST(
root=root, # 数据存储路径
train=False, # 加载测试集（共10000个样本）
transform=transform # 应用相同的数据变换
)
# 数据加载器配置参数
# 根据是否使用GPU选择不同的优化参数
# 将num_workers设为0避免多进程问题
kwargs = {'num_workers': 0, 'pin_memory': True} if use_cuda else {'num_workers': 0}
# 创建训练数据加载器
train_loader = DataLoader(
dataset=train_set, # 训练数据集
batch_size=batch_size, # 每个批次的样本数
shuffle=True, # 每个epoch打乱数据顺序，防止模型记忆顺序
**kwargs # 解包上述配置参数
)
# 创建测试数据加载器
test_loader = DataLoader(
dataset=test_set, # 测试数据集
batch_size=batch_size, # 批次大小（通常与训练集相同）
shuffle=False, # 测试集不需要打乱，保证可重复性
**kwargs # 解包配置参数
)
return train_loader, test_loader
4. 模型训练#
该函数实现了量子变分自编码器（Q-VAE）在 MNIST 数据集上的完整训练流程。 模型结合了经典神经网络编码器/解码器与受限玻尔兹曼机（RBM），通过最小化带权重衰减的负 ELBO 损失进行优化，并引入 KL 散度控制潜在分布与先验的对齐程度。 训练过程中记录各项损失指标并定期保存至文件。
def train(self, run_tsne=False, compute_energy=False, tsne_interval=10, generate_animation=False):
"""执行完整训练流程"""
logger.info(f"Start training {self.config.type}")
self._setup_data()
self._create_model()
self._setup_tuner()
# 提取测试集标签（用于能量可视化）
y_test = []
for _, labels in self.test_loader:
y_test.extend(labels.numpy().tolist())
y_test = np.array(y_test)
# 调试信息
logger.info(f"Test labels - unique: {np.unique(y_test)}, count: {len(y_test)}")
logger.info(f"Test labels distribution: {np.bincount(y_test)}")
tsne_frames = []
epoch_pbar = tqdm(range(1, self.config.num_epochs + 1), desc="Training Progress")
for epoch in epoch_pbar:
train_loss = self.tuner.train(epoch)
if hasattr(train_loss, 'item'):
train_loss = train_loss.item()
self.train_losses.append(train_loss)
test_loss, input_data, output_data, label_list = self.tuner.test()
if hasattr(test_loss, 'item'):
test_loss = test_loss.item()
self.test_losses.append(test_loss)
epoch_pbar.set_description(f"Epoch {epoch}/{self.config.num_epochs} - Train Loss: {train_loss:.2f}, Test Loss: {test_loss:.2f}")
if epoch % (self.config.num_epochs // 10) == 0 or epoch == self.config.num_epochs:
self._save_reconstruction(epoch, input_data, output_data)
logger.info(f"Epoch {epoch}: Train Loss={train_loss:.4f}, Test Loss={test_loss:.4f}")
if generate_animation and (epoch % tsne_interval == 0 or epoch == self.config.num_epochs):
frame_path = self._save_tsne_frame(epoch)
tsne_frames.append(frame_path)
self.tuner.save_model(config_string=f"final_{self.config.type}")
self._plot_training_curve()
if generate_animation and tsne_frames:
create_tsne_animation(tsne_frames, output_path=self.output_dir)
import shutil
frame_dir = os.path.join(self.output_dir, "temp_tsne_frames")
if os.path.exists(frame_dir):
shutil.rmtree(frame_dir)
if run_tsne:
logger.info("Generating t-SNE visualization...")
self._visualize_tsne()
if compute_energy:
logger.info("Compute BM energy for each sample...")
energies = self.compute_energy(self.model, self.test_loader)
self._save_energy_plot(energies, y_test, output_path=self.output_dir)
logger.info(f"{self.config.type} training completed")
return self.model, self.train_losses, self.test_losses
5. 可视化与评估#
本节提供两类关键可视化工具：
一是通过 plot_training_curves 绘制训练/验证损失与准确率曲线，用于监控模型收敛情况；
二是利用 t_SNE 对 QVAE 模型提取的潜在表示进行降维可视化，揭示不同类别在隐空间中的分布结构。
两者均支持自动保存高分辨率图像，并可灵活控制是否实时显示，便于实验分析、结果记录。
5.1 训练过程可视化#
def plot_training_curves(
train_loss_history,
val_loss_history,
train_acc_history,
val_acc_history,
save_path=None,
show=True,
):
"""
绘制训练和验证的损失及准确率曲线
Args:
train_loss_history: 训练损失历史
val_loss_history: 验证损失历史
train_acc_history: 训练准确率历史
val_acc_history: 验证准确率历史
save_path: 图像保存路径
"""
plt.figure(figsize=(12, 5))
# 损失曲线
plt.subplot(1, 2, 1)
plt.plot(train_loss_history, label="Training Loss", color="blue", alpha=0.7)
plt.plot(val_loss_history, label="Validation Loss", color="red", alpha=0.7)
plt.title("Training and Validation Loss")
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.legend()
plt.grid(True, alpha=0.3)
# 准确率曲线
plt.subplot(1, 2, 2)
plt.plot(train_acc_history, label="Training Accuracy", color="blue", alpha=0.7)
plt.plot(val_acc_history, label="Validation Accuracy", color="red", alpha=0.7)
plt.title("Training and Validation Accuracy")
plt.xlabel("Epoch")
plt.ylabel("Accuracy (%)")
plt.legend()
plt.grid(True, alpha=0.3)
plt.tight_layout()
# plt.close()
# 自动保存
if save_path is None:
# 生成默认保存路径
# timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
save_path = f"results/mlp_training_curves_{timestamp}.png"
plt.savefig(save_path, dpi=300, bbox_inches="tight")
print(f"Training curves saved to: {save_path}")
plt.show()
if show:
plt.show()
else:
plt.close() # 不显示时关闭图像，节省内存
5.2 潜在空间可视化#
def t_SNE(
test_loader,
qvae_model,
point_size=20,
alpha=0.6,
epochs=None,
save_path=None,
show=True,
):
"""
QVAE版本的t-SNE可视化
Args:
test_loader: 测试数据加载器
qvae_model: QVAE模型
use_std: 是否使用标准差
point_size: 点大小
alpha: 透明度
epochs: 训练轮数
save_path: 保存路径
show: 是否显示图像
"""
features = []
labels = []
qvae_model.eval()
device = next(qvae_model.parameters()).device
with torch.no_grad():
for batch_idx, (example_data, example_targets) in enumerate(test_loader):
example_data = example_data.to(device)
# QVAE前向传播 - 获取潜变量zeta
_, _, _, zeta = qvae_model(example_data)
zeta_np = zeta.cpu().numpy()
for idx in range(zeta_np.shape[0]):
features.append(zeta_np[idx])
labels.append(example_targets[idx].item())
# 创建DataFrame
feat_cols = [f"dim_{i}" for i in range(zeta_np.shape[1])]
df = pd.DataFrame(features, columns=feat_cols)
df["label"] = labels
df["label"] = df["label"].apply(lambda i: str(i))
logger.info(f"Extracted {len(features)} samples with {zeta_np.shape[1]} dimensions")
# 执行t-SNE
logger.info("Running t-SNE...")
tsne = TSNE(n_components=2, verbose=1, perplexity=30, max_iter=500, random_state=42)
tsne_results = tsne.fit_transform(df[feat_cols].values)
df_tsne = df.copy()
df_tsne["x-tsne"] = tsne_results[:, 0]
df_tsne["y-tsne"] = tsne_results[:, 1]
# 可视化
fig, ax = plt.subplots(figsize=(10, 8))
scatter = ax.scatter(
df_tsne["x-tsne"],
df_tsne["y-tsne"],
c=df_tsne["label"].astype(int),
cmap="tab10",
s=point_size,
alpha=alpha,
)
# 添加颜色条
cbar = plt.colorbar(scatter, ax=ax, label="Digit")
# 动态标题和文件名
# training_status = "fully_trained" if epochs and epochs >= 20 else f"epochs_{epochs}"
training_status = f"epochs_{epochs}"
title = f"t-SNE Visualization of QVAE Latent Space ({training_status})"
plt.title(title)
plt.xlabel("t-SNE dimension 1")
plt.ylabel("t-SNE dimension 2")
plt.tight_layout()
# 自动保存
if save_path is None:
# 生成默认保存路径
# timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
save_path = f"results/t-SNE_QVAE_{training_status}_{timestamp}.png"
plt.savefig(save_path, dpi=300, bbox_inches="tight")
logger.info(f"t-SNE plot saved to: {save_path}")
plt.show()
if show:
plt.show()
else:
plt.close() # 不显示时关闭图像，节省内存
return df_tsne, save_path, training_status
6. 表征学习与分类#
Q-VAE 学到的表征可用于下游分类任务：
该函数 train_mlp_classifier 用于训练一个多层感知机(MLP)分类器，输入特征是通过QVAE模型提取的数据表征。
它首先将数据集划分为训练集和验证集，并初始化一个MLP模型、优化器和损失函数。
在每个训练周期，模型参数根据训练集更新，并在验证集上评估性能。
def fit(self, X, y, validation_split=0.2):
"""训练MLP模型"""
if self.input_dim is None:
self.input_dim = X.shape[1]
logger.info(f"Auto-detected input_dim: {self.input_dim}")
# 数据划分
X_train, X_val, y_train, y_val = train_test_split(
X, y,
test_size=validation_split,
random_state=self.random_state,
stratify=y
)
# 转为 Tensor
X_train = torch.FloatTensor(X_train).to(self.device)
y_train = torch.LongTensor(y_train).to(self.device)
X_val = torch.FloatTensor(X_val).to(self.device)
y_val = torch.LongTensor(y_val).to(self.device)
self.classes_ = np.unique(y)
# 创建数据加载器
train_dataset = TensorDataset(X_train, y_train)
val_dataset = TensorDataset(X_val, y_val)
train_loader = DataLoader(train_dataset, batch_size=self.batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=self.batch_size, shuffle=False)
# 创建模型
self.model = self._create_model()
optimizer = torch.optim.Adam(self.model.parameters(), lr=self.lr, weight_decay=self.weight_decay)
criterion = nn.CrossEntropyLoss()
# 记录训练历史
train_loss_history = []
val_loss_history = []
train_acc_history = []
val_acc_history = []
best_val_acc = 0.0
best_state = None
epoch_pbar = tqdm(range(1, self.epochs + 1), desc="Training MLP")
for epoch in epoch_pbar:
# 训练
train_acc, avg_train_loss = self._train_mlp_epoch(
model=self.model,
data_loader=train_loader,
optimizer=optimizer,
criterion=criterion,
device=self.device
)
# 验证
val_acc, avg_val_loss = self._eval_mlp_epoch(
model=self.model,
data_loader=val_loader,
criterion=criterion,
device=self.device
)
# 记录历史
train_loss_history.append(avg_train_loss)
val_loss_history.append(avg_val_loss)
train_acc_history.append(train_acc)
val_acc_history.append(val_acc)
# 打印进度
if epoch % 10 == 0:
# print(f"Epoch {epoch}: Train Loss={train_loss:.4f}, Val Loss={val_loss:.4f}, Val Acc={val_acc:.2f}%")
logger.info(
f"Epoch {epoch}/{epoch}: "
f"Train Loss={avg_train_loss:.4f}, Train Acc={train_acc:.2f}% | "
f"Val Loss={avg_val_loss:.4f}, Val Acc={val_acc:.2f}%"
)
# 选择最佳模型
if val_acc > best_val_acc:
best_val_acc = val_acc
best_state = self.model.state_dict()
if self.save_path is not None:
model_save_path = os.path.join(self.save_path, "best_mlp_classifier.pth")
torch.save(best_state, model_save_path)
# 加载最佳模型
self.model.load_state_dict(best_state)
logger.info(f"Best Validation Accuracy: {best_val_acc:.2f}%")
# 绘制训练曲线
if self.save_path is not None:
curves_save_path = os.path.join(
self.save_path, f"mlp_training_curves_epochs_{self.epochs}.png"
)
else:
curves_save_path = None # let plot_training_curves generate default
plot_training_curves(
train_loss_history=train_loss_history,
val_loss_history=val_loss_history,
train_acc_history=train_acc_history,
val_acc_history=val_acc_history,
save_path=curves_save_path,
show=True,
)
return self
7. 科研应用：QBM-VAE#
Q-VAE 的进阶版本 QBM-VAE 在科研中展示了重要价值：
单细胞转录组学分析：
显著提升聚类精度
检测传统方法无法辨识的新型细胞亚型
为靶点发现提供新线索