RBM Classification: Handwritten Digit Recognition#
This tutorial demonstrates how to use a Restricted Boltzmann Machine (RBM) for feature learning and classification on the handwritten-digit dataset. It is suitable for beginners to understand the workflow of applying RBM to image feature extraction and classification.
Tutorial Objectives#
After finishing this tutorial you will know how to:
Use an RBM to extract features from image data
Combine an RBM feature extractor with a traditional classifier
Visualize the weights learned by the RBM and the samples it generates
Evaluate the performance of the classification model
Runtime Environment#
Example location: example/rbm_digits/rbm_digits.ipynb
Dependencies:
pip install scikit-learn matplotlib scipy
1. Data Preparation#
This section shows how to load the classic handwritten-digit dataset and expand it with simple data-augmentation techniques (e.g., shifting images in four directions) to improve the model’s generalization ability.
First, we use load_digits from sklearn to load the data, obtaining 8×8-pixel handwritten-digit images and their labels. Next, each original image is shifted up, down, left, and right to create additional training samples. Finally, all images are flattened into a 2-D array, split into training and test sets, and normalized so that every feature value lies in the interval [0, 1], readying the data for subsequent model training.
def load_data(self, plot_img=False):
"载入图片数据"
digits = load_digits()
images = digits.images # 8x8 的图像矩阵
labels = digits.target # 对应的标签
# 获取图像数据和标签
# 扩展数据集
expanded_images = []
expanded_labels = []
for image, label in zip(images, labels):
# 原始图像
expanded_images.append(image)
expanded_labels.append(label)
# 向四个方向平移
for direction in ["up", "down", "left", "right"]:
translated_image = self.translate_image(image, direction)
expanded_images.append(translated_image)
expanded_labels.append(label)
# 将列表转换为 NumPy 数组
expanded_images = np.array(expanded_images)
expanded_labels = np.array(expanded_labels)
# 可视化图像数据和标签
if plot_img:
plt.figure(figsize=(16, 9))
for index in range(5):
plt.subplot(1, 5, index + 1)
plt.imshow(expanded_images[index], origin="lower", cmap="gray")
plt.title("Training: %i\n" % expanded_labels[index], fontsize=18)
# 将图像数据展平为二维数组 (n_samples, 64)
n_samples = expanded_images.shape[0]
data = expanded_images.reshape((n_samples, -1))
# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(
data, expanded_labels, test_size=0.2, random_state=42
)
# 使用sklearn的MinMaxScaler进行归一化
scaler = MinMaxScaler()
X_train = scaler.fit_transform(X_train)
X_test = scaler.transform(X_test)
return X_train, X_test, y_train, y_test
2. Model Training#
This section shows how to train a Restricted Boltzmann Machine (RBM) model.
First, the RBM model is initialized according to the input feature dimension and the chosen number of hidden units, and the model is placed on the desired device (CPU or GPU). Then, a stochastic-gradient-descent (SGD) optimizer is defined to update the model parameters. The entire training process uses mini-batch gradient descent to gradually adjust parameters so as to minimize an objective function that includes the negative log-likelihood plus weight- and bias-decay terms.
If verbose mode is enabled, the current iteration’s objective value and parameter statistics are printed periodically, and generated sample images together with model parameters are visualized, making it easy to monitor learning progress and effectiveness in real time.
def fit(self, X, y=None): # 修改接口以符合scikit-learn约定
"""
训练RBM模型
Args:
X: 训练数据，形状为 (n_samples, n_features)
y: 忽略，为兼容scikit-learn接口
"""
# 初始化受限玻尔兹曼机（RBM）模型
rbm = RestrictedBoltzmannMachine(
X.shape[1], # 可见层单元数（特征维度）
self.n_components, # 隐层单元数
)
rbm.to(self.device) # 将模型移动到指定设备（CPU/GPU）
self.rbm = rbm
# 初始化优化器
opt_rbm = SGD(rbm.parameters(), lr=self.learning_rate)
n_samples = X.shape[0] # 样本数量
n_batches = int(np.ceil(float(n_samples) / self.batch_size)) # 批次数量
# 生成每个batch的切片索引
batch_slices = list(
gen_even_slices(n_batches * self.batch_size, n_batches, n_samples=n_samples)
)
X_torch = torch.FloatTensor(X).to(self.device) # 转为torch张量并移动到设备
idx = 0
# 训练循环
for iteration in range(1, self.n_iter + 1):
for step, batch_slice in enumerate(batch_slices):
idx += 1
x = X_torch[batch_slice] # 获取当前batch数据
x = rbm.get_hidden(x) # 正相（计算隐层激活）
s = rbm.sample(self.sampler) # 负相（采样重构数据）
# s = rbm.get_visible(x[:, rbm.num_visible :]) # 使用隐藏层重构可见层
opt_rbm.zero_grad() # 梯度清零
# 计算目标函数（等价于负对数似然），并加权衰减项
w_weight_decay = 0.02 * torch.sum(rbm.quadratic_coef**2) # 权重衰减
b_weight_decay = 0.05 * torch.sum(rbm.linear_bias**2) # 偏置衰减
objective = rbm.objective(x, s) + w_weight_decay + b_weight_decay
# 反向传播并更新参数
objective.backward()
opt_rbm.step()
# 如果verbose，定期评估模型性能和可视化参数
if self.verbose:
print(f"Iteration {idx}, Objective: {objective.item():.6f}")
if (idx - 1) % 20 == 0:
# 打印权重和偏置的均值与最大值
print(
f"jmean {torch.abs(rbm.quadratic_coef).mean()}"
f" jmax {torch.abs(rbm.quadratic_coef).max()}"
)
print(
f"hmean {torch.abs(rbm.linear_bias).mean()}"
f" hmax {torch.abs(rbm.linear_bias).max()}"
)
if self.plot_img:
display_samples = (
rbm.sample(self.sampler)
.cpu()
.numpy()[:20, : rbm.num_visible]
)
# 生成样本
plt.figure(figsize=(16, 2))
plt.imshow(self.gen_digits_image(display_samples, 8))
plt.title(f"Generated samples at iteration {iteration}")
plt.show()
_, axes = plt.subplots(1, 2)
axes[0].imshow(rbm.quadratic_coef.detach().cpu().numpy())
axes[1].imshow(
rbm.quadratic_coef.grad.detach().cpu().numpy()
)
plt.tight_layout()
plt.show()
return self
3. Feature Extraction and Classification#
3.1 Training the Classifier#
This subsection shows how to build and train a two-stage classification pipeline consisting of an RBM followed by logistic regression, and compares its performance with a logistic-regression model that uses only the raw pixel features.
First, the handwritten-digit dataset is loaded and pre-processed via RBMRunner. The RBM is then embedded as a feature extractor inside a Pipeline, followed by a logistic-regression classifier. Training time for both approaches is recorded, and accuracy together with a classification report are evaluated on the test set.
The experiment aims to verify whether the high-level features learned by the RBM can improve the performance of a downstream classifier, and also provides an intuitive example for understanding the role of unsupervised pre-training in deep learning.
def train_classifier(n_iter=2, use_cim=False):
logistic = LogisticRegression(random_state=42)
# 初始化RBM
rbm = RBMRunner(
n_components=128,
learning_rate=0.1,
batch_size=32,
n_iter=n_iter,
verbose=True,
plot_img=False,
random_state=seed,
use_cim=use_cim,
)
# 加载数据
X_train, X_test, y_train, y_test = rbm.load_data(plot_img=True)
classifier = Pipeline(steps=[("rbm", rbm), ("logistic", logistic)])
########## 训练模型 ##########
logistic.C = 500.0
logistic.max_iter = 1000
# 训练 RBM-Logistic Pipeline
start_time = time.time()
classifier.fit(X_train, y_train)
training_time = time.time() - start_time
print(f"RBM Pipline training completed in {training_time:.2f} seconds")
# 训练 Logistic regression
logistic_classifier = LogisticRegression(C=500.0, max_iter=1000, random_state=42)
start_time = time.time()
logistic_classifier.fit(X_train, y_train)
training_time = time.time() - start_time
print(f"Logistic regression training completed in {training_time:.2f} seconds")
########## 评估模型 ##########
pip_pred = classifier.predict(X_test)
pip_acc = accuracy_score(y_test, pip_pred)
print(
"\nLogistic regression using RBM features:\n%s\n"
% (classification_report(y_test, pip_pred))
)
print(f"Test Accuracy: {pip_acc:.4f}")
log_pred = logistic_classifier.predict(X_test)
log_acc = accuracy_score(y_test, log_pred)
print(
"\nLogistic regression using raw pixel features:\n%s\n"
% (classification_report(y_test, log_pred))
)
print(f"Test Accuracy: {log_acc:.4f}")
return rbm, y_test, log_pred
3.2 Visualizing Weights#
This function visualizes the weights associated with each hidden unit of the RBM as 8×8 images, giving an intuitive view of the feature patterns learned from the data. The result can be saved as a high-resolution PDF.
def plot_weights(self, save_as="qbm_weights", save_pdf=False):
"""绘制权重"""
weights = self.rbm.quadratic_coef.detach().cpu().numpy()
fig, axes = plt.subplots(
8, 16, gridspec_kw={"wspace": 0.1, "hspace": 0.1}, figsize=(16, 7)
)
fig.suptitle(f"{self.n_components} components extracted by QBM", fontsize=16)
fig.subplots_adjust()
for i, ax in enumerate(axes.flatten()):
if i < weights.shape[1]:
ax.imshow(weights[:, i].reshape(8, 8), cmap=plt.cm.gray)
ax.axis("off")
# 保存结果
if save_pdf:
_ensure_result_dir()
plt.savefig(
f"results/{save_as}.pdf", dpi=300, bbox_inches="tight", format="pdf"
)
plt.show()
3.2 Visualizing the Confusion Matrix#
This function visualizes the model’s confusion matrix on the test set as a heat-map, clearly showing prediction accuracy and confusion among classes. A custom suffix can be added to the title, and the figure can be saved as a high-resolution PDF.
def plot_confusion_matrix(self, y_true, y_pred, title_suffix="", save_pdf=False):
"""绘制混淆矩阵"""
cm = confusion_matrix(y_true, y_pred)
plt.figure(figsize=(10, 8))
sns.heatmap(cm, annot=True, fmt="d", cmap="Blues")
plt.title(f"Confusion Matrix ({title_suffix})", fontsize=18)
plt.xlabel("Predicted Label", fontsize=16)
plt.ylabel("True Label", fontsize=16)
# plt.xticks(rotation=45)
# 保存结果
if save_pdf:
_ensure_result_dir()
plt.savefig(
f"results/rbm_confusion_matrix_{title_suffix}.pdf",
dpi=300,
bbox_inches="tight",
format="pdf",
)
plt.tight_layout()
plt.show()