https://github.com/fxia22/stn.pytorch

pytorch version of spatial transformer networks
https://github.com/fxia22/stn.pytorch

pytorch stn

Last synced: 5 days ago
JSON representation

pytorch version of spatial transformer networks

• Host: GitHub
• URL: https://github.com/fxia22/stn.pytorch
• Owner: fxia22
• License: other
• Created: 2017-02-23T06:39:35.000Z (almost 8 years ago)
• Default Branch: master
• Last Pushed: 2017-06-21T23:01:33.000Z (over 7 years ago)
• Last Synced: 2024-12-29T22:26:20.309Z (12 days ago)
• Topics: pytorch, stn
• Language: Jupyter Notebook
• Size: 15.7 MB
• Stars: 589
• Watchers: 16
• Forks: 87
• Open Issues: 16
• Metadata Files:
  • Readme: README.md
  • License: LICENSE

Awesome Lists containing this project

README

        
# PyTorch version of spatial transformer network

Ported from https://github.com/qassemoquab/stnbhwd according to pytorch tutorial. Now support CPU and GPU. To use the ffi you need to install the `cffi` package from pip.

# Build and test

```

cd script

./make.sh #build cuda code, don't forget to modify -arch argument for your GPU computational capacity version

python build.py

python test.py

```

There is a demo in `test_stn.ipynb`

# Modules

`STN` is the spatial transformer module, it takes a `B*H*W*D` tensor and a `B*H*W*2` grid normalized to [-1,1] as an input and do bilinear sampling.

`AffineGridGen` takes a `B*2*3` matrix and generate an affine transformation grid. 

`CylinderGridGen` takes a `B*1` theta vector and generate a transformation grid to remap equirectangular images along x axis. 

`DenseAffineGridGen` takes a `B*H*W*6` tensor and do affine transformation for each pixel. Example of convolutional spatial transformer can be found in `test_conv_stn.ipynb`.

An example of the landscape of the loss function of a simple STN with L1 Loss can be found in the demo.      

# Train hacks

- set a learning rate multiplier, 1e-3 or 1e-4 would work fine. 

- add an auxiliary loss to regularized the difference of the affine transformation from identity mapping, to aviod sampling outside the original image. 

# Complex grid demo

STN is able to handle a complex grid, however, how to parameterize the grid is a problem.

![image](stn.jpg)