Pytorch Implementation of RefineNet
This is a Pytorch implementation of the Multipath RefineNet architecture from the paper.
Installation
Install PyTorch following instructions on their website. Then install this package:
pip install git+https://github.com/thomasjpfan/pytorch_refinenet.git
Implemented versions
- Multi-path 4-Cascaded RefineNet:
RefineNet4Cascade
- Multi-path 4-Cascaded RefineNet With Improved Pooling:
RefineNet4CascadePoolingImproved
There are diagrams of these two versions in the author's github repo. The improved pooling version adds an additional pooling/convolution layer and flips the order of the pooling/convolution layers in the Chained Residual Pooling block.
Usage
This implementation of the Multipath RefineNet has the following initialization:
class RefineNet4Cascade(nn.Module):
def __init__(self,
input_shape,
num_classes=1,
features=256,
resnet_factory=models.resnet101,
pretrained=True,
freeze_resnet=True):
...
The input_shape
is a tuple of(channels, size)
which denotes the number of channels in the
input image and the input width/height. For an input to flow cleanly through the resnet layers, the input size should be divisible by 32. The input size is assumed to be a square image/patch. For
example the RefineNet4Cascade
can be defined to intake 3x224x224 images:
import torch
from pytorch_refinenet import RefineNet4Cascade
net = RefineNet4Cascade((3, 224), num_classes=10)
x = torch.randn(1, 3, 224, 224)
y = net(x)
y.size()
# torch.Size([1, 10, 56, 56])
The number of channels outputed will equal num_classes
and the size will be 1/4 the size of the
input as described in the paper. You can upscale the to get back to the original resolution.
Training
The refinenet backbone is frozen by default, which means they will not be updated with gradients during training.
The parameters
method in RefineNet4Cascade
was redefined to only return the parameters that require a gradident. Thus this will work for training:
net = RefineNet4Cascade((3, 224), num_classes=10)
opt = optim.Adam(net.parameters())
x = torch.randn(1, 3, 224, 224)
y = net(x)
...