pytorch中roipooling_Pytorch的那些坑

pytorch是最火也是最方便的deep learning框架（科研用途），但是在使用过程中也遇到了好多坑，尤其版本更新的时候。因此建立一个帖子，持续记录pytorch使用过程中遇到的坑。如果大家遇到其他bug或者对分析有更深的见解，欢迎补充，我会merge更新

1. 多GPU并行
2. roipooling, roialign, nms

多GPUs并行

”RuntimeError: Expexted tensor for argument #1 'input' to have the same device as tensor for argument #2 'weight'; but device 0 does not equal 2(while checking argument for cudnn_convolution)

我实在pytorch0.31和1.6上遇到的这个问题，其他版本有没有未知。欢迎report。

这个问题是由于用于计算的网络module和输入的数据不在同一个device上导致的，比如上面这个，网络的’weight'在gpu2上，但是数据‘input'在gpu1上。最后半句’ but device 0 does not equal 2‘我没弄懂是什么原因，pytorch官方社区说是gpu0造成的，但是实验表明和pytorch机制表明并不是。sdf使用比如CUDA_VISIBLE_DEVICES=1,3,4后， pytorch和tf都会对gpu从0开始再编号。经过2天debug终于找到原因了。我原来网络的定义如下：

def __init__(self,.....):
  if use_resnet:
    self.features, self.layer4 = load_resnet()
            self.roi_fmap = self._roi_map
  else: # vgg
   vgg = load_vgg()
    self.features = vgg.features
    self.roi_map = vgg.classifiers

  def forward(self, x):
     ...
     return self.roi_fmap(x)sdfsdfsdfsdf
  ....
  # for convenient, I wrap the operation here
  def _roi_map(self, features):
        return self.layer4(features).mean(3).mean(2) 

我使用vgg的时候多gpu没有问题，所以可以排除是gpu0造成的。debug的时候发现每次都是_roi_map 这一步报错。这个是别人的faster rcnn的pyroch版本写法，我本身就很不喜欢_roi_map的这种跳来跳去的定义，因此我修改为

def __init__(self,.....):
  if use_resnet:
    self.features, self.layer4 = load_resnet()
            self.roi_fmap = self.layer4
  else: # vgg
   vgg = load_vgg()
    self.features = vgg.features
    self.roi_map = vgg.classifiers

  def forward(self, x):
     ...
     out = self.roi_fmap(x)
     return out.mean(3).mean(2)
  ....

然后问题解决了。

我的分析是， 用function _roi_map（）把self.layer4(features).mean(3).mean(2) wrap起来导致在model初始化的时候，pytorch没有发现网络的这一步操作，从而没有将它分配到各个gpu上，但是在forward的时候调用了这一步，所以在默认gpu上进行了操作，也就是gpu0，从而导致input和weight的device不匹配。debug的时候的确也证明，第一步的时候input在0的时候，没有报错，而是到其它gpu的时候开始报错（input.device可以知道当前参数在那个gpu）。

我猜测，定义网络的时候，网络每个layer或者module或者model的定义必须直接定义在继承nn.Model类的__init__（）下，pytorch才能在一开始将其分配到各个gpu（通过eplicas = nn.parallel.replicate(model, devices=list(range(num_gpus)))函数）。pytorch不会顺着往下找各个定义在类下面的function是否还有其他layer或者module。上面例子中，如果把_roi_map定义为一个类:

class _roi_map(nn.Model):
  def __init__(self,layer4):
    self.roi_map = layer4.mean(3).mean(2)
  def forward(self,x):
    return self.roi_map(x)

.....
....
def __init__(self,.....):
  if use_resnet:
    self.features, self.layer4 = load_resnet()
            self.roi_fmap = _roi_map(layer4)

同样可以解决这个问题。

roipooling, roialign, nms

这三个操作在two-stage object detection框架必不可少。之前都是写c文件然后编译成cuda使用。但是pytorch1.0后不再支持c文件编译了，而网上可找到的这三个文件的c++写法很少。另外，不同的机器，cuda版本都要重新编辑，非常麻烦，还经常出问题。pytorhc1.6提供了这三个算法的官方函数：

from torchvision.ops import nms
from torchvision.ops import RoIAlign, RoIPool

# 而且scores 和boxes 还不需要人为的实现排序好，而之前c文件编写的普遍需要先排序才能输入nms
keep = nms(boxes, scores, iou_threshold)

# spatial_scale就是backbone最终输出的feature map相对于原图输入缩放的倍数，一般是1/16 或者1/32
# sampling_ratio可以不用设置， 提出roalign的原paper设置为2
roi_align = RoIAlign((pooling_heigh,pooling_width), spatial_scale, sampling_ratio)
roi_feat = roi_align(feat_map, boxes)

#RoIPool使用方式类似, 只是没有sampling_ratio
roi_pool = RoIPool((pooling_heigh,pooling_width), spatial_scale)
roi_feat = roi_pool (feat_map, boxes)

使用官方的函数在visual genome和coco上测试性能和速度并没有下降