FM's Devlog | Debugging a Skin Disease Classification Notebook

NB: this one assumes you know the basics of Keras and tf.data, I will not be writing a glossary this time

Prefix

I came across a machine learning notebook that had two problems. It was using around 47 gigs of RAM during training, which kept getting close to crashing the runtime, and accuracy was stuck around 71%. I decided to look into why.

The notebook is a dual scope dermatological image classification pipeline. Scope A is 10 broad skin disease classes, Scope B is 4 classes that look visually similar to each other (Eczema, Psoriasis, Seborrh Keratoses, Tinea), meant to stress test the models on a harder, more ambiguous problem. For each scope it trains two model families, a custom CNN (a static baseline and a Hyperband tuned version) and EfficientNetB3 (a frozen baseline and a fine tuned version). So 8 models total, evaluated with classification reports, confusion matrices, and t-SNE plots of the learned feature space.

The notebook ran top to bottom and produced numbers, so on the surface nothing looked broken. But once I actually read through it line by line, I found a handful of things that were either silently wrong, silently wasteful, or just fragile. None of these would throw an error. They would just quietly give worse models or burn more RAM than needed. That is the annoying part about ML pipelines, a bug doesn’t crash, it just makes your accuracy a bit lower and you never know.

So here’s a walkthrough of what I changed and, more importantly, why. By the end the RAM usage never crossed 13 gigs.

Bug 1: EfficientNetB3 was getting 224x224 images

This is the one that stood out the most once I noticed it. The data pipeline resized every image to (224, 224) because that’s what the custom CNN expects. EfficientNetB3’s canonical input resolution is (300, 300). It was being fed the wrong size the entire time.

It still runs, because Keras will happily resize internally or just work with a smaller input than the network was designed around, but you lose the resolution the pretrained weights were tuned against. So I split the pipeline into two parallel tracks, one at 224x224 for the custom CNN and one at 300x300 specifically for EfficientNetB3:

IMG_SIZE = (224, 224)
IMG_SIZE_EFF = (300, 300)

and built two sets of datasets off the same directories, just with different target_size passed into build_pipeline. Small fix, but it’s the kind of thing that would have quietly capped how well the transfer learning models could ever do, and is likely part of why accuracy was sitting at 71%.

Bug 2: caching float32 images was the main RAM problem

This was the bigger one. The original pipeline rescaled pixels to [0, 1] as float32 and then called .cache() on that. That means every image sits in memory as 4 bytes per channel. Across train, valid, and test for both scopes, at full resolution that adds up to the 47 gigs of RAM that was causing crashes.

The fix was to cache the images as uint8 (1 byte per channel, native [0, 255] range) and only cast to float32 after pulling from the cache:

ds = ds.map(lambda x, y: (tf.cast(x, tf.uint8), y), num_parallel_calls=AUTOTUNE)
ds = ds.cache(cache_path) if cache_path else ds.cache()
ds = ds.map(lambda x, y: (tf.cast(x, tf.float32), y), num_parallel_calls=AUTOTUNE)

That alone brought memory down from 47 gigs to under 13 gigs across the whole run, with zero accuracy cost, since the actual scaling still happens before the model sees anything. The other thing to be careful about was order. Resizing has to happen before .cache(), not after, otherwise you’re caching full resolution images and the memory savings disappear. Caching at native resolution is apparently a classic way to accidentally blow up your RAM, and this notebook was doing exactly that.

Bug 3: the “divide then multiply back” rescaling dance

The old version rescaled every image to [0, 1] in the pipeline, then for EfficientNetB3 added an inverse Rescaling(255.0) layer right at the model input to undo it, because EfficientNet expects raw [0, 255] pixels. That round trip works, but it’s fragile, and it’s the kind of thing that’s easy to forget about and get backwards later.

So now the pipeline just keeps images in their native [0, 255] range, and each model handles its own scaling internally. The custom CNN gets a Rescaling(1./255) layer at its input, EfficientNetB3 just takes the raw pixels since that’s already what it expects. One source of truth for “what range does this model want,” instead of two pipelines fighting each other.

Bug 4: class weights weren’t used during the hyperparameter search

Balanced class weights were being computed to fight the class imbalance, and passed into the final model.fit() calls. But the Hyperband tuner.search() calls weren’t getting class_weight at all. That means the tuner was picking the “best” architecture under one set of conditions (no imbalance correction), and then the actual final model was trained under a different set of conditions (with imbalance correction). The model selection and the final training weren’t even optimizing the same thing.

Fixed by just passing class_weight=class_weights into tuner.search() as well, so the search and the final fit agree on what “good” means.

Bug 5: BatchNorm layers getting unfrozen during fine tuning

This one probably also contributed to the accuracy ceiling. When unfreezing the top layers of EfficientNetB3 for fine tuning, the code was unfreezing everything in that range, including the BatchNormalization layers. Updating BatchNorm running statistics on a small medical imaging dataset is a known way to destabilize a pretrained backbone, since those stats were computed over ImageNet’s scale of data, not a few thousand skin lesion photos.

The fix:

base_backbone.trainable = True
for layer in base_backbone.layers[:-40]:
    layer.trainable = False
for layer in base_backbone.layers:
    if isinstance(layer, layers.BatchNormalization):
        layer.trainable = False

Unfreeze the top block, then explicitly re-freeze any BatchNorm layer inside that block. This is apparently the recommended Keras pattern for fine tuning and the notebook just wasn’t doing it.

Bug 6: jumping straight into fine tuning with a fresh head

Related to the above, the original code went straight from “tuner found the best classification head” to “unfreeze backbone and train at 1e-5.” But that classification head was randomly initialized right before this step. Trying to train a random head and a barely thawed backbone at the same time, at an extremely low learning rate, doesn’t give the head enough room to actually learn anything useful before the low learning rate kicks in.

So I added a warmup phase. First train the new head with the backbone still fully frozen at a normal learning rate, then unfreeze and fine tune at the low rate:

print("\n>>> Transfer Phase 2a: Warming up classification head (backbone frozen)...")
tuned_model.fit(train_ds, validation_data=valid_ds, epochs=8, ...)

print("\n>>> Transfer Phase 2b: Unfreezing top backbone layers for specialization...")
# unfreeze, refreeze BatchNorm, recompile at 1e-5, fit again

Two short phases instead of one long shot in the dark.

Smaller things I cleaned up along the way

A few other things that aren’t worth their own section but added up:

• Determinism: separate tf.random.set_seed(42) and np.random.seed(42) calls were replaced with a single keras.utils.set_random_seed(42) call, which seeds TensorFlow, NumPy, and Python’s own random module together.
• Mixed precision: added keras.mixed_precision.set_global_policy('mixed_float16') when a GPU is present. Roughly halves memory usage and speeds things up on T4 or better. The output layer of every model now explicitly keeps dtype='float32', since softmax under float16 can get numerically unstable.
• Label smoothing: switched the loss to CategoricalCrossentropy(label_smoothing=0.05). Clinical labels aren’t perfectly clean, a small amount of smoothing helps calibration.
• Tuner objective: switched Hyperband from val_accuracy to val_loss. Accuracy ties too easily on small validation sets and doesn’t reflect confidence, loss is a steadier signal for picking the best trial.
• EarlyStopping and ReduceLROnPlateau everywhere: previously some training calls had early stopping and some didn’t, and the patience values were inconsistent. Pulled all of this into one make_callbacks() helper so every model trains under the same stopping and learning rate decay rules.
• Dataset copying: the original code physically copied every image into a new folder structure for each scope. Now it tries a symlink first and only falls back to copying if symlinks aren’t supported. Saves disk space and is much faster to set up.
• Split name resolution: hardcoded 'train', 'valid', 'test' folder names everywhere. Some datasets call it 'val' or 'validation' instead, so I added a small resolver that checks for known aliases instead of assuming.
• Fail fast on missing classes: added a validation step right after building the scope folders that checks every expected class exists in every split, and raises immediately with a clear message if not. Better than burning GPU hours into a run that crashes three models in because one folder name doesn’t match.
• t-SNE embedding layer: previously it grabbed the embedding from a hardcoded layer index, model.layers[-2]. That breaks the moment you change the architecture even slightly. Now it searches for the actual GlobalAveragePooling2D layer by type, which is the layer the embedding should be coming from anyway.
• try/finally around training tracks: the per-track memory cleanup (release_memory) was only called after a successful run. Wrapped each track in try/finally so memory still gets released even if a particular model run throws partway through.

Conclusion

None of these bugs would have shown up as an error message. The notebook ran fine before, it just wasn’t doing what it looked like it was doing in a few places. EfficientNetB3 was working with the wrong input size, the tuner was optimizing for a different objective than the final training, and the fine tuning setup was quietly working against the pretrained weights instead of with them. On top of that, caching full resolution float32 images was the main reason RAM usage was hitting 47 gigs. The lesson here is that “it runs and gives a number” is a pretty low bar for an ML pipeline. The real check is going back through every step and asking whether each one is doing exactly what it’s supposed to be doing, not just whether it throws.

Postfix

This was a good exercise in slowing down and reading through a notebook line by line instead of just trusting that it worked because the metrics looked reasonable. Getting RAM usage from 47 gigs down to under 13, while also fixing the input size mismatch and the BatchNorm freezing, should move accuracy well past that 71% starting point. I’m curious to see the actual numbers once both versions are compared side by side. Stay tuned for that one.