DEV Community

Sowaiba Arshad
Sowaiba Arshad

Posted on

Why My Medical AI Took 6.4 Seconds Per Scan and How I Got It to 3.1.

I built a chest X-ray diagnostic platform called ThoraxNet. It detects 14 thoracic pathologies from a single image, reports how confident it is using Monte Carlo Dropout, draws **GradCAM **heatmaps over the regions that drove each prediction, and writes a structured radiology report with an LLM.

The model was the interesting part to build. It was not the part that decided whether the thing felt like a product. That came down to one unglamorous number: how long a user waits after they hit "analyze."

This post is about how I took that wait from 6.4 seconds to 3.1 seconds, why the honest answer is "about 2x and not 10x," and the two bugs I only found because I stopped trusting the happy path and read the logs.

Live demo: https://thorax-tho.vercel.app
Code: https://github.com/Sowaiba-01/ThoraxNet

The problem: a good model wrapped in a slow request

The model worked. The demo worked. But every scan took roughly six and a half seconds, and a six second spinner makes anything feel broken no matter how good the output underneath it is.

My first instinct was the same wrong instinct most people have: "the vision transformer forward pass must be the bottleneck, so let me go optimize the model." I have shipped enough regressions chasing that kind of hunch to know it is worth nothing until it is measured.

So before changing a single line of model code, I instrumented the pipeline to report how long each stage actually took, and returned that breakdown on every single response:

"stage_timings_ms": {
  "preprocess": 16,
  "mc_dropout": 2868,
  "gradcam": 6,
  "report": 92
}
Enter fullscreen mode Exit fullscreen mode

The result was not what I expected, and that is exactly why measuring first matters. A single forward pass through the transformer is about 15 milliseconds. The model was never the problem. Monte Carlo Dropout was, and it was eating nearly all of the server side budget.

The root cause: twenty forward passes, one at a time

Monte Carlo Dropout estimates uncertainty by running the model many times with dropout left switched on, then measuring how much the predictions move. My original implementation did the obvious thing:

samples = []
for _ in range(n_samples):        # n_samples = 20
    logits = model(x)             # one image, batch size 1
    samples.append(torch.sigmoid(logits))
Enter fullscreen mode Exit fullscreen mode

Twenty sequential forward passes, each with a batch size of one. Every pass pays the full per call overhead, and the hardware spends most of its time waiting between launches instead of computing.

The fix is to stop asking twenty separate times and ask once. Tile the single image into a batch of twenty and run one forward pass:

tiled = x.repeat(n_samples, 1, 1, 1)   # (20, 3, 224, 224)
logits = model(tiled)                  # ONE forward pass
probs = torch.sigmoid(logits).view(n_samples, batch, -1)
mean, std = probs.mean(0), probs.std(0)
Enter fullscreen mode Exit fullscreen mode

The correctness argument is the part worth understanding, and it is the exact question a good interviewer will ask you. Are those twenty copies not identical? No. Dropout samples a fresh mask for every element in the batch. So the twenty tiled copies each get an independent dropout mask, which means they are exactly the twenty independent stochastic samples the estimator needs. Same statistics, one launch instead of twenty.

I did not want to trust that reasoning on faith, so I wrote a test that runs both the old sequential version and the new batched version four hundred times each and asserts their means agree within Monte Carlo error:

assert torch.allclose(batched_mean, sequential_mean, atol=0.05)
Enter fullscreen mode Exit fullscreen mode

That test is the difference between "I think this is equivalent" and "I proved this is equivalent."

The results, measured on identical hardware

Same host, same image, same protocol of thirty requests with warmup discarded. The only variable is the code.

Metric Before (v1.0.0) After (v1.1.0) Improvement
p50 latency 6,387 ms 3,127 ms 1.9x faster
p95 latency 7,525 ms 3,862 ms 1.9x faster
p99 latency 8,026 ms 4,768 ms 1.7x faster
Throughput 0.15 req/s 0.30 req/s 2.0x

Every percentile improved by roughly 2x, with no change to the model weights and no change to accuracy. This is purely a change in how the model is executed.

Two smaller changes cleaned up the rest of the request path. The LLM report call was moved off the critical path onto a worker thread, so the user gets their pathology result without waiting on report generation. And GradCAM heatmaps, which were being recomputed on every request, are now cached per image and class.

The honest part: why it is 2x and not 10x

Here is where it would be easy to lie, and where lying would eventually cost me an offer.

On a GPU, this exact change is often close to 10x. The entire win comes from keeping an accelerator busy that was otherwise sitting idle between twenty tiny kernel launches.

ThoraxNet runs on a free tier CPU host with two virtual cores. There is far less idle parallelism to reclaim there. So the same code change gives roughly 2x, not roughly 10x.

I could have written "10x faster" in the title and most readers would never have checked. But the number that survives a technical interview is the one you can explain. The honest version is: 2x on CPU, because the batching win is bounded by how much idle parallelism there is to reclaim, and Monte Carlo Dropout on CPU is still the dominant cost, so the next real lever is GPU inference or quantization, not more batching. That sentence is worth more than a bigger number I cannot defend.

Two bugs I found by reading the whole path, not the ticket

Profiling forced me to read the entire request path from HTTP call to response. Two problems fell out that had nothing to do with latency.

GradCAM had never worked. The route handler read an attribute on the pipeline that was never actually assigned anywhere. The inference code built the heatmaps into a local variable and then dropped them when the function returned. Every heatmap request came back as a 404. Nothing logged an error, and the frontend simply rendered an empty panel, so it shipped and sat broken because no code path ever raised. I fixed it and pinned it with a regression test that fails if the overlays are not recorded, because a silent bug earns a loud test.

Every radiology report was quietly failing. While watching the deploy logs I saw this on repeat:

The model `llama3-70b-8192` has been decommissioned and is no longer supported.
Enter fullscreen mode Exit fullscreen mode

The LLM provider had retired that model months earlier. Every report request was returning a 400 and falling back to a template, so users were getting canned text instead of a generated report, and nobody noticed because the fallback made it look fine. The fix was one line to point at the current model, plus an environment variable so the next deprecation is a config change and not a code change.

Neither of these was in my task list. Both were only visible because I stopped trusting the happy path and actually read what the server was doing.

What I would do next

Monte Carlo Dropout still dominates at roughly 2.8 seconds. Batching took the easy win. The remaining levers are honest about the hardware rather than clever in the request path:

GPU or quantized inference to cut the per request cost of the twenty passes. The export and evaluation scripts are already written, and the accuracy delta table is the homework I still owe.

A shared cache for the GradCAM session store, since it currently lives in process memory that will not survive a restart or a second replica.

The takeaway

The model was 15 milliseconds. The product was 6.4 seconds. The entire gap lived in how the model was called, not in the model itself, and I only found that because I measured before I touched anything and read the logs instead of the ticket.

The unglamorous work is the work. Profiling first, proving equivalence with a test, and being honest about a 2x instead of inflating it to a 10x are the habits that hold up when someone actually reads your code.


If you found this useful, follow me on GitHub for more write ups like this: https://github.com/Sowaiba-01

The full project is open source. Fork it, build on it, and if it helped you or you just think it is cool, a star means a lot: https://github.com/Sowaiba-01/ThoraxNet

ThoraxNet is for research use only. It is not FDA cleared and is not a substitute for a radiologist.

Top comments (1)

Collapse
 
ahmetozel profile image
Ahmet Özel

"The model was 15 milliseconds, the product was 6.4 seconds" is the sentence I wish more ML posts started with. Almost every latency problem I have chased lived in serialization, per-request preprocessing, or a synchronous call nobody batched - never in the forward pass. Also worth flagging your Monte Carlo Dropout number: 20 stochastic passes is a product decision disguised as a modelling one. If the uncertainty estimate is not changing what happens downstream, it is 2.8 seconds bought for nothing.