maskrcnn-benchmark 代码详解之　modeling/utils.py

高效torch.cat函数

最新推荐文章于 2020-04-08 21:28:40 发布

原创最新推荐文章于 2020-04-08 21:28:40 发布 · 719 阅读

0 ·

CC 4.0 BY-SA版权

文章标签：

#maskrcnn_benchmark #utils.py #cat #目标检测 #RCNN

maskrcnn benchmark 专栏收录该内容

25 篇文章

订阅专栏

本文介绍了一个重写的cat函数，该函数在输入为单个张量时避免了不必要的复制操作，仅在存在多个张量需要拼接时执行拼接，提高了效率。

def cat(tensors, dim=0):
    """
    Efficient version of torch.cat that avoids a copy if there is only a single element in a list
    """
    """
    重写了cat函数，使得当传入的张量只有一个元素的时候避免做拼接操作，只有可拼接的情况下，即有多个元素的情况下进行拼接
    """
    assert isinstance(tensors, (list, tuple))
    if len(tensors) == 1:
        return tensors[0]
    return torch.cat(tensors, dim)