Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
A snippet to show how roialign works
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import numpy as np
def pyth_crop_pool_layer(bottom, rois, pool_size):
x1 = rois[:, 1::4] / 16.0
y1 = rois[:, 2::4] / 16.0
x2 = rois[:, 3::4] / 16.0
y2 = rois[:, 4::4] / 16.0
height = bottom.size(2)
width = bottom.size(3)
# affine theta
zero = rois.data.new(rois.size(0), 1).zero_()
theta = torch.cat([\
(x2 - x1) / (width - 1),
zero,
(x1 + x2 - width + 1) / (width - 1),
zero,
(y2 - y1) / (height - 1),
(y1 + y2 - height + 1) / (height - 1)], 1).view(-1, 2, 3)
grid = F.affine_grid(theta, torch.Size((rois.size(0), 1, pool_size, pool_size)))
crops = F.grid_sample(bottom.expand(rois.size(0), bottom.size(1), bottom.size(2), bottom.size(3)), grid)
return crops
bottom = Variable(torch.arange(0,49).view(1,1,7,7))
rois = Variable(torch.Tensor([[0,0,0,6,6], [0,0,1,3,6]])) * 16 # ignore rois[:, 0]
pool_size = 4
print(pyth_crop_pool_layer(bottom, rois, pool_size))
print(bottom[0:7,0:7])
@simo23
Copy link

simo23 commented Dec 13, 2017

Thanks for the gist!

For sure it could be a bug on my configuration, but to make this work I had to change line 23 to:

    zero = Variable(torch.FloatTensor([[0],[0]]))   # [2x1 FloatTensor]

torch.cat was complaining about having a list as 0 argument.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment