[bugfix] - support large CPU-side datasets; fix numpy/CuPy interop

train.py: add _to_gpu() helper; convert mini-batches from CPU numpy to GPU
CuPy on the fly so large datasets can stay in RAM. Use numpy permutation for
numpy batches during shuffle. Limit status err_rms to first 5000 samples to
avoid GPU OOM on status checks.

model.py: replace np.copy() with np.array() — cupy.copy() does not convert
numpy arrays to CuPy; cupy.array() does. Fixes entity.forward() crash after
training when batch was passed as numpy.

test_faces_sub_image.py: load_patches normalizes on CPU and returns numpy,
avoiding multi-copy GPU OOM for large datasets (500 images ≈ 1.4 GB).
show_reconstructions converts numpy patches to CuPy before inference.
Add --n_images CLI arg.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
2026-05-30 22:26:06 +02:00
co-authored by Claude Sonnet 4.6
parent 8c10bda002
commit 2f9fe1b66b
3 changed files with 30 additions and 10 deletions
+14 -4
View File
@@ -77,7 +77,11 @@ def load_patches(data_dir: str, n_images: int = N_IMAGES, start: int = 0):
patches_np = numpy.concatenate(all_patches, axis=0)
std = numpy.std(patches_np, axis=1)
patches_np = patches_np[std > 0.01]
return normalize(np.array(patches_np))
# Normalize on CPU; return numpy — train() converts mini-batches to GPU on the fly
mean = patches_np.mean(axis=1, keepdims=True)
std2 = patches_np.std(axis=1, keepdims=True)
return (patches_np - mean) / numpy.where(std2 < 1e-8, 1.0, std2)
def load_image_patches(path: str):
@@ -130,13 +134,17 @@ def show_filters(entity: Entity, rows: int = 8, cols: int = 16):
plt.tight_layout()
def show_reconstructions(model: TestModel, patches: Mat, n_show: int = 10):
def show_reconstructions(model: TestModel, patches, n_show: int = 10):
fig, axes = plt.subplots(2, n_show, figsize=(n_show * 1.5, 3.5))
axes[0, 0].set_ylabel('Original')
axes[1, 0].set_ylabel('Recon')
fig.suptitle('Patch reconstructions (normalised display)', fontsize=10)
err_total = 0.0
# Convert to GPU if patches arrived as a CPU numpy array
if isinstance(patches, numpy.ndarray):
patches = np.array(patches)
cmap = 'gray' if GRAYSCALE else None
def to_img(p):
@@ -183,6 +191,8 @@ if __name__ == '__main__':
help='Number of training epochs (default: 1000)')
ap.add_argument('--mini_batch_size', type=int, default=1000,
help='Mini-batch size (default: 1000)')
ap.add_argument('--n_images', type=int, default=N_IMAGES,
help=f'Number of training images (default: {N_IMAGES})')
args = ap.parse_args()
if args.grayscale:
@@ -201,8 +211,8 @@ if __name__ == '__main__':
model.load()
if args.do_train:
print(f'Loading {N_IMAGES} training images (stride={STRIDE})...')
train_patches = load_patches(DATA_DIR, N_IMAGES)
print(f'Loading {args.n_images} training images (stride={STRIDE})...')
train_patches = load_patches(DATA_DIR, args.n_images)
print(f'Training on {train_patches.shape[0]} patches ({N_VIS}-dim each)')
model.train(train_patches)
model.save()