[refactor] move tests to tests/, add pytest functions and main() entrypoints
- Moved src/tests/ → tests/ - Added test_* functions with assertions to script-style test files - Added main() to each so IDEs offer it as a separate run target from pytest - Fixed cupy_test.py: remove spurious x_gpu += x_cpu, fix duplicate xlabel Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,86 @@
|
||||
import random
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
from model.model import Model
|
||||
from rbm.entity import Entity, EntityParams, TrainingParams
|
||||
from rbm.matrix import Mat, np
|
||||
|
||||
N = 8
|
||||
N_VIS = N * N
|
||||
N_HID = 32
|
||||
N_CASES = 400
|
||||
|
||||
|
||||
class BinaryAEModel(Model):
|
||||
def __init__(self, name: str, work_dir: str = '.'):
|
||||
super().__init__(name, work_dir)
|
||||
self.unit1 = Entity(
|
||||
(N_VIS, N_HID),
|
||||
EntityParams(),
|
||||
TrainingParams(learning_rate=0.1, momentum=0.9, num_epochs=5000, do_rao_blackwell=True)
|
||||
)
|
||||
|
||||
def forward(self, x: Mat) -> Mat:
|
||||
return self.unit1.forward(x)
|
||||
|
||||
def reconstruct(self, x: Mat) -> Mat:
|
||||
return self.unit1.reconstruct(x)
|
||||
|
||||
|
||||
def make_bars_and_stripes(n_cases: int, n: int = N) -> Mat:
|
||||
data = np.zeros((n_cases, n * n), dtype=np.float64)
|
||||
for i in range(n_cases):
|
||||
img = np.zeros((n * n,), dtype=np.float64)
|
||||
if random.random() > 0.5:
|
||||
row = random.randint(0, n - 1)
|
||||
img[row * n:(row + 1) * n] = 1.0
|
||||
else:
|
||||
col = random.randint(0, n - 1)
|
||||
img[col::n] = 1.0
|
||||
data[i] = img
|
||||
return data
|
||||
|
||||
|
||||
def test_binary_autoencoder():
|
||||
model = BinaryAEModel("binary_autoencoder_pytest", "results")
|
||||
model.unit1.training_params.num_epochs = 500
|
||||
model.init(0.01)
|
||||
train_batch = make_bars_and_stripes(N_CASES)
|
||||
model.train(train_batch)
|
||||
mae = float(np.mean(np.abs(train_batch - np.array([model.reconstruct(model.forward(x)) for x in train_batch]))))
|
||||
assert mae < 0.4, f"Reconstruction MAE {mae:.4f} too high"
|
||||
|
||||
|
||||
def main():
|
||||
model = BinaryAEModel("binary_autoencoder", "results")
|
||||
model.init(0.01)
|
||||
|
||||
train_batch = make_bars_and_stripes(N_CASES)
|
||||
model.train(train_batch)
|
||||
model.save()
|
||||
|
||||
test_batch = make_bars_and_stripes(10)
|
||||
n_show = len(test_batch)
|
||||
|
||||
fig, axes = plt.subplots(2, n_show, figsize=(n_show * 1.5, 3))
|
||||
axes[0, 0].set_ylabel("Original")
|
||||
axes[1, 0].set_ylabel("Recon")
|
||||
|
||||
total_error = 0.0
|
||||
for i, inp in enumerate(test_batch):
|
||||
h = model.forward(inp)
|
||||
recon = model.reconstruct(h)
|
||||
total_error += float(np.mean(np.abs(inp - recon)))
|
||||
|
||||
axes[0, i].imshow(np.asnumpy(np.reshape(inp, (N, N))), cmap='gray', vmin=0, vmax=1)
|
||||
axes[0, i].axis('off')
|
||||
axes[1, i].imshow(np.asnumpy(np.reshape(recon, (N, N))), cmap='gray', vmin=0, vmax=1)
|
||||
axes[1, i].axis('off')
|
||||
|
||||
print(f"Mean reconstruction error: {total_error / n_show:.4f}")
|
||||
plt.tight_layout()
|
||||
plt.show()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user