An ENet implementation in PyTorch, based on convert_torch_to_pytorch
Switch branches/tags
Nothing to show
Clone or download
Pull request Compare This branch is 8 commits ahead of clcarwin:master.
Fetching latest commit…
Cannot retrieve the latest commit at this time.
Type Name Latest commit message Commit time
Failed to load latest commit information.


Fork based on convert_torch_to_pytorch and adapted for image segmentation with ENet in PyTorch.

Conversion from torch ENet

  • Convert the source torch7 trained model from cutorch/cudnn weights to CPU weights for conversion:
$ th convert_torch_gpu_to_cpu.lua -model
  • Convert the lua model to a pytorch script and state dict:
$ python --model --output enet_pytorch

This produces a standalone script defining a torch.nn module, and an associated state dict.

Example usage:

import torch
from torch.autograd import Variable
import torch.nn.functional as F
import cv2
from enet_pytorch import enet_pytorch


img = 'frankfurt_000001_030067_leftImg8bit.png'
im = cv2.imread(img).astype(np.float32)[:, :, ::-1]/255
inp = torch.from_numpy(im.transpose(2, 0, 1)).unsqueeze(0).cuda()
inp = F.upsample(Variable(inp), (512, 1024), mode='bilinear').data

out = enet_pytorch(Variable(inp))

Tested to reproduce ENet original results.

Model and weights

For reference the output on the author's cityscapes model and associated converted weights are provided on The specifies the conversion of training IDs to original cityscapes dataset IDs for this model.

Scores on the validation set

The author's converted model's scores on cityscapes validation set are as follows

Class IoU Class iIoU Category IoU Category iIoU
60.4 36.2 80.0 63.9

Implementation notes

The max-unpooling modules unpool the last pooling module (FIFO), which is appropriate for encoder-decoder segmentation networks. This converter is not tested for other networks than ENet (but should remain compatible with the models tested with convert_torch_to_pytorch).



The converter code developed in this repository follows the MIT license. ENet's model and weights (provided externally to this repository) are part of ENet and follow the CC BY-NC 4.0 license.