如果你想使用带有CUDA的Opencv dnn和带有GPU（可选）的torch，我建议使用此方法：
FROM nvidia/cuda:10.2-base-ubuntu18.04
WORKDIR /home
ENV DEBIAN_FRONTEND=noninteractive
ENV TZ=Europe/Minsk
RUN ln -snf /usr/share/zoneinfo/$TZ /etc/localtime && echo $TZ > /etc/timezone
RUN apt-get update && apt-get install -y \
keyboard-configuration \
nvidia-driver-440\
curl \
ca-certificates \
sudo \
git \
bzip2 \
libx11-6 \
cmake \
g++ \
wget \
build-essential \
cmake \
git \
unzip \
pkg-config \
python-dev \
python-opencv \
libopencv-dev \
libjpeg-dev \
libpng-dev \
libtiff-dev \
libgtk2.0-dev \
python-numpy \
python-pycurl \
libatlas-base-dev \
gfortran \
webp \
python-opencv \
qt5-default \
libvtk6-dev \
zlib1g-dev \
libcudnn7=7.6.5.32-1+cuda10.2 \
libcudnn7-dev=7.6.5.32-1+cuda10.2 \
python3-pip \
python3-venv \
nano
RUN alias python='/usr/bin/python3'
RUN pip3 install numpy
RUN pip3 install torch
RUN git clone https://github.com/opencv/opencv_contrib && \
cd opencv_contrib && \
git fetch --all --tags && \
git checkout tags/4.3.0 && \
cd .. && \
git clone https://github.com/opencv/opencv.git && \
cd opencv && \
git checkout tags/4.3.0
RUN cd opencv && mkdir build && cd build && \
cmake -DCMAKE_BUILD_TYPE=Release \
-D CMAKE_CXX_COMPILER=/usr/bin/g++ \
-D PYTHON_DEFAULT_EXECUTABLE=$(which python3) \
-D BUILD_NEW_PYTHON_SUPPORT=ON \
-D BUILD_opencv_python3=ON \
-D HAVE_opencv_python3=ON \
-D INSTALL_PYTHON_EXAMPLES=ON \
-D CUDA_TOOLKIT_ROOT_DIR=/usr/local/cuda-10.2 \
-D CUDA_BIN_PATH=/usr/local/cuda-10.2 \
-D CUDNN_INCLUDE_DIR=/usr/include/cudnn.h \
-D WITH_CUDNN=ON \
-D CUDA_ARCH_BIN=6.1 \
-D OPENCV_DNN_CUDA=ON \
-D WITH_CUDA=ON \
-D BUILD_opencv_cudacodec=OFF \
-D WITH_GTK=ON \
-D CMAKE_BUILD_TYPE=RELEASE \
-D CUDA_HOST_COMPILER:FILEPATH=/usr/bin/gcc-7 \
-D ENABLE_PRECOMPILED_HEADERS=OFF \
-D WITH_TBB=ON \
-D WITH_OPENMP=ON \
-D WITH_IPP=ON \
-D BUILD_EXAMPLES=OFF \
-D BUILD_DOCS=OFF \
-D BUILD_PERF_TESTS=OFF \
-D BUILD_TESTS=OFF \
-D WITH_CSTRIPES=ON \
-D OPENCV_EXTRA_MODULES_PATH=../../opencv_contrib/modules \
-D CMAKE_INSTALL_PREFIX=/usr/local/ \
-DBUILD_opencv_python3=ON \
-D PYTHON_DEFAULT_EXECUTABLE=$(which python3) \
-D PYTHON3_EXECUTABLE=$(which python3) \
-D PYTHON_INCLUDE_DIR=$(python3 -c "from distutils.sysconfig import get_python_inc; print(get_python_inc())") \
-D PYTHON3_INCLUDE_DIR=$(python3 -c "from distutils.sysconfig import get_python_inc; print(get_python_inc())") \
-D PYTHON3_LIBRARY=$(python3 -c "from distutils.sysconfig import get_config_var;from os.path import dirname,join ; print(join(dirname(get_config_var('LIBPC')),get_config_var('LDLIBRARY')))") \
-D PYTHON3_NUMPY_INCLUDE_DIRS=$(python3 -c "import numpy; print(numpy.get_include())") \
-D PYTHON3_PACKAGES_PATH=$(python3 -c "from distutils.sysconfig import get_python_lib; print(get_python_lib())") \
-D OPENCV_GENERATE_PKGCONFIG=ON .. \
-Dopencv_dnn_superres=ON /usr/bin/ .. && \
make -j$(nproc) && \
make install
RUN pip3 install opencv/build/python_loader
然后你可以运行
import torch
import os
print('availabe:',torch.cuda.is_available() )
print('devices available', torch.cuda.device_count())
print('device id:',torch.cuda.current_device() )
print('device address', torch.cuda.device(0))
print('gpu model',torch.cuda.get_device_name(0))
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('Using device:', device)
if device.type == 'cuda':
print(torch.cuda.get_device_name(0))
print('Memory Usage:')
print('Allocated:', round(torch.cuda.memory_allocated(0)/1024**3,1), 'GB')
print('Cached: ', round(torch.cuda.memory_reserved(0)/1024**3,1), 'GB')
import cv2
print("DNN_BACKEND_CUDA",cv2.dnn.DNN_BACKEND_CUDA)
print("DNN_BACKEND_CUDA",cv2.dnn.DNN_TARGET_CUDA)
然后你会得到类似这样的东西：
Using device: cuda
availabe: True
devices available 1
device id: 0
device address <torch.cuda.device object at 0x7f5a0a392550>
gpu model GeForce GTX 1050 Ti
Memory Usage:
Allocated: 0.0 GB
Cached: 0.0 GB
DNN_BACKEND_CUDA 5
DNN_BACKEND_CUDA 6
