mapbox / robosat

Semantic segmentation on aerial and satellite imagery. Extracts features such as: buildings, parking lots, roads, water, clouds
MIT License
2.02k stars 383 forks source link

Works around internal PyTorch bug related to shared weights #96

Closed daniel-j-h closed 6 years ago

daniel-j-h commented 6 years ago

The PyTorch tracing model exporter has an issue with shared weights / refs to layers, see

https://github.com/pytorch/pytorch/issues/8392

The workaround is to not store any refs to the resnet layers in the model itself. Instead we have to directly pull the input tensors through the resnet layers in the forward pass.

This will change the .pth file format. We could work around that by loading the serialized state dict in non-strict fashion in the export tool. But then we are opening up the door to letting actual errors going through.

Note: we need to fix the FPN pull request in a similar way https://github.com/mapbox/robosat/pull/75

cc @bkowshik @maning

daniel-j-h commented 6 years ago

This needs a release with major version bump since we are breaking the .pth format.