Hello everyone, hope you are all having a great day.
I was wondering if its possible to convert the following module into torch script:
class PriorBox(torch.nn.Module):
def __init__(self):
super(PriorBox, self).__init__()
# @torch.jit.script
def forward(self, minSizes, steps, clip, image_size):
anchors = []
feature_maps = [[ceil(image_size[0]/step), ceil(image_size[1]/step)] for step in steps]
for k, f in enumerate(feature_maps):
min_sizes = minSizes[k]
for i, j in product(range(f[0]), range(f[1])):
for min_size in min_sizes:
s_kx = min_size / image_size[1]
s_ky = min_size / image_size[0]
dense_cx = [x * steps[k] / image_size[1] for x in [j + 0.5]]
dense_cy = [y * steps[k] / image_size[0] for y in [i + 0.5]]
for cy, cx in product(dense_cy, dense_cx):
anchors += [cx, cy, s_kx, s_ky]
# back to torch land
output = torch.Tensor(anchors).view(-1, 4)
if clip:
output.clamp_(max=1, min=0)
return output
I have reimplemented this in libtorch but its slow as snail!
here it is in case its needed:
struct PriorBox : torch::nn::Module
{
PriorBox(const std::vector<std::pair<torch::Tensor, torch::Tensor>>& min_sizes, const std::vector<int>& steps, const bool& clip, const std::pair<int, int>& image_size, const std::string& phase = "train")
{
this->min_sizes = min_sizes;
this->steps = steps;
this->clip = clip;
this->image_size = image_size;
this->phase = phase;
this->name = name;
for (auto& step : this->steps)
{
auto height = torch::tensor(float(this->image_size.first) / step);
auto width = torch::tensor(float(this->image_size.second) / step);
this->feature_maps.emplace_back(std::make_pair(torch::ceil(height), torch::ceil(width)));
}
}
torch::Tensor forward()
{
std::vector<torch::Tensor> anchors;
int i = -1;
for (auto& fmap : this->feature_maps)
{
auto min_sizes = this->min_sizes[++i];
auto result = torch::cartesian_prod({ torch::arange(c10::Scalar(0), c10::Scalar(fmap.first.item().toInt()), c10::Scalar(1)),
torch::arange(c10::Scalar(0), c10::Scalar(fmap.second.item().toInt()), c10::Scalar(1)) });
for (int idx = 0; idx <= result.sizes()[0] - 1; idx++)
{
//takes around 0.006 ms
auto i_ = result[idx][0];
auto j_ = result[idx][1];
//takes around 0.20 ms to 0.30 ms
for (auto& min_size : { min_sizes.first, min_sizes.second })
{
auto s_kx = min_size / float(this->image_size.second);
auto s_ky = min_size / float(this->image_size.first);
// takes around 0.037 ms
torch::Tensor dense_cx = (j_ + 0.5) * this->steps[i] / this->image_size.second;
torch::Tensor dense_cy = (i_ + 0.5) * this->steps[i] / this->image_size.first;
//takes around 0.02ms
auto result_cy_cx = torch::cartesian_prod({ dense_cy.unsqueeze(0), dense_cx.unsqueeze(0) });
//this takes around 0.010ms
for (int l = 0; l <= result_cy_cx.sizes()[0] - 1; l++)
{
auto cy = result_cy_cx[l][0].unsqueeze(0);
auto cx = result_cy_cx[l][1].unsqueeze(0);
anchors.emplace_back(torch::cat({ cx, cy, s_kx, s_ky }));
}
}
}
}
//takes around 5ms!
auto output = torch::stack(anchors).view({ -1,4 });
//takes around 0 ms!
if (this->clip)
output.clamp_(0, 1);
return output;
}
std::vector<std::pair<torch::Tensor, torch::Tensor>> min_sizes;
std::vector<int> steps;
bool clip;
std::pair<int, int> image_size;
std::string phase;
std::string name;
std::vector<std::pair<torch::Tensor, torch::Tensor>> feature_maps;
};
So I thought maybe converting this into a Torchscript and loading it in C++ would be a better idea.
Since we are dealing with loops, etc we cant simply jit trace the module. So I guess I’m stuck with scripting.
When I tried to do scripting I kept getting different errors such asrange cannot be used as a value:
RuntimeError:
range cannot be used as a value:
File "P:\ligen\layers\functions\prior_box.py", line 19
for k, f in enumerate(feature_maps):
min_sizes = minSizes[k]
for i, j in product(range(int(f[0])), range(int(f[1]))):
~~~~~~~~~~~~~~ <--- HERE
for min_size in min_sizes:
s_kx = min_size / image_size[1]
So my question is, considering this, is it portable to torchscript? if so What am I missing here? how should I go about this?
If not, what are my other options?
Thanks a lot in advance