[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:
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user