Google Colab无法看到.so文件

Question

我正在尝试进行CUDA扩展以在Google Colab中的Python中使用它。

这是setup.py

sources = [os.path.join(path_to_cpp, 'interface.cpp'),
            os.path.join(path_to_cpp, 'box_encoder_cuda.cu'),
            os.path.join(path_to_cpp, 'random_horiz_flip.cu')]

setup(
    name='SSD',
    ext_modules=[
        CUDAExtension(
            'SSD._C', 
            sources)
    ],
    cmdclass={
        'build_ext': BuildExtension
    })

构建时，我可以清楚地看到libc10.so存在于构建中

x86_64-linux-gnu-g++ -pthread -shared -Wl,-O1 -Wl,-Bsymbolic-functions -Wl,-Bsymbolic-functions -Wl,- 
z,relro -Wl,-Bsymbolic-functions -Wl,-z,relro -g -fstack-protector-strong -Wformat -Werror=format- 
security -Wdate-time -D_FORTIFY_SOURCE=2 build/temp.linux-x86_64-3.6/content/drive/My  
Drive/SSD/csrc/interface.o build/temp.linux-x86_64-3.6/content/drive/My Drive/SSD/csrc/box_encoder_cuda.o build/temp.linux-x86_64-3.6/content/drive/My  
Drive/SSD/csrc/random_horiz_flip.o -L/usr/local/lib/python3.6/dist-packages/torch/lib - 
L/usr/local/cuda/lib64 -lc10 -ltorch -ltorch_cpu -ltorch_python -lcudart -lc10_cuda -ltorch_cuda -o  
build/lib.linux-x86_64-3.6/SSD/_C.cpython-36m-x86_64-linux-gnu.so  

但是，如果我尝试在Colab中使用它：

%%writefile test.py
import SSD._C

print("All ok")

!python test.py

我将看到以下错误消息：

Traceback (most recent call last):
  File "test.py", line 1, in <module>
    import SSD._C
ImportError: libc10.so: cannot open shared object file: No such file or directory

也许是因为没有在Python运行时中声明库，所以我必须以某种方式声明它们吗？预先谢谢！

P.S。当我尝试直接在单元格中使用它而不写入.py文件时：

import SSD._C

print("All ok")

产量

ModuleNotFoundError: No module named 'SSD'

Answer 1

在导入扩展名之前通过导入炬管解决的问题：

%%writefile test.py
import torch
import SSD._C

print("All ok")

根据documentation“这将解决动态链接器必须看到的一些符号”。

但是，为什么不能在.ipynb单元格中直接使用此扩展名仍然是一个开放点。

编辑：如果在.ipynb中安装新扩展名，请确保在使用运行时之前重新启动它。重新启动后，您将可以在.ipynb单元格中使用新的扩展名