Weight Conversion Pitfalls: MXNet to PyTorch

The optimistic version of an MXNet-to-PyTorch port goes like this: export the Gluon model's parameters, map them onto an equivalent PyTorch module, run an input through both, see similar numbers, ship it. Sometimes that works. The dangerous part is that when it doesn't work, it usually almost works — the converted model loads cleanly, produces plausible outputs, and diverges from the original in ways a smoke test won't catch.

These are the pitfalls we hit most often in real conversions, roughly in order of how quietly they fail.

BatchNorm semantics

Batch normalization is the classic silent killer. Three separate traps live here:

  • Epsilon placement and defaults. MXNet's BatchNorm uses eps=1e-5 by default, as does PyTorch — but models frequently override it, and a conversion that maps weights while dropping a non-default epsilon produces outputs that differ in the third or fourth decimal place. Small enough to survive a glance; large enough to move decisions near a classification boundary.
  • Momentum conventions. MXNet's momentum (default 0.9) weights the existing running statistic; PyTorch's momentum (default 0.1) weights the new observation. They are complements of each other. This doesn't matter for frozen inference, but the moment someone fine-tunes the ported model with naively copied hyperparameters, the running statistics drift differently than the original's did.
  • use_global_stats and training mode. Gluon models sometimes run BatchNorm in global-stats mode even during training. A port that reproduces the architecture but not this flag trains differently and can even infer differently if evaluation-mode handling isn't faithful.

Padding is not padding

MXNet and PyTorch both let you say padding=1, and for most convolutions the results agree. The divergence shows up with asymmetric padding: pooling layers with ceil_mode behavior, and architectures ported originally from other frameworks (TensorFlow's SAME padding especially) that baked asymmetric pads into the MXNet implementation. MXNet's pooling defaults and PyTorch's differ in when they apply ceil versus floor to output sizes; a one-pixel difference in a feature map early in the network is not a one-pixel problem by the final layer.

Check every pooling layer's output shape against the original, not just the network's final output shape — shape agreement at the end can mask crops and pads that shifted content spatially.

Layout: NCHW, NHWC, and what the weights think

Gluon computer-vision models are typically NCHW, matching PyTorch — but not always, and exported symbolic models sometimes carry layout transposes inside the graph. Convolution weights transferred without honoring a layout difference produce garbage obvious enough to catch; dense layers after a flatten are the subtle case. Flattening NCHW versus NHWC activations orders the features differently, so a fully-connected layer's weight matrix must be permuted accordingly. Outputs will be wrong but structured — often still producing a plausible-looking probability distribution.

Custom operators and hybridized control flow

HybridBlock code that branches on input shapes, custom operators written against the MXNet C++ API, and SymbolBlock models loaded from JSON have no mechanical translation. The porting choice is re-implementation, and re-implementation means the original's quirks — including its bugs — must be reproduced or consciously fixed. A custom op with a subtly nonstandard gradient doesn't matter for inference parity but matters enormously if you plan to fine-tune.

Preprocessing lives outside the model

Half the "conversion bugs" we diagnose are not in the model at all. GluonCV's default image pipeline (resize interpolation method, crop behavior, normalization constants, channel order) differs in small ways from torchvision's defaults. Bilinear interpolation alone differs between implementations — OpenCV, PIL, and MXNet's image module do not produce identical resized images. If the original service used MXNet's image decoding and the new one uses PIL, you can achieve perfect model parity and still ship different predictions.

Port the preprocessing with the same rigor as the network, and test parity end-to-end from raw input, not from preprocessed tensors.

Numerical environment

The same model, faithfully converted, can still diverge via the execution environment: cuDNN autotuning selecting different convolution algorithms, TF32 being enabled by default on Ampere-class GPUs in PyTorch, fused versus unfused operations reordering floating-point arithmetic. These differences are usually within tolerance — but you need a defined tolerance to know that, measured on representative inputs, with TF32 and autotuning settings pinned during comparison.

What protects you

Every pitfall above shares a property: invisible to "it runs," visible to systematic comparison. The protection is a parity harness — identical inputs through original and port, outputs compared at documented tolerances, on a sample that covers the input distribution's edges, with layer-wise comparison available for localizing divergence. Build it before the port, not after; it turns each of these pitfalls from a production incident into a failing test.

We've written up our parity methodology separately — and when a conversion fails parity persistently, the honest answer is sometimes retraining rather than archaeology. Knowing when to stop converting and start retraining is half the craft.