Can’t import 'ConvBNReLU' from 'torchvision.models.mobilenet' while using GPU-CodePudding

Problem

I’m following this tutorial to train an object detection model for Lens Studio. I want to use a GPU to train the model, either Google Colab’s cloud GPU or my local through Jupyter.

When I try to import:

import torchvision
from torchvision.models.mobilenet import ConvBNReLU
from torchvision.ops import box_iou, nms

I'm getting the following error when the runtime is set to GPU.

---------------------------------------------------------------------------
ImportError                               Traceback (most recent call last)
<ipython-input-3-157e58f8c700> in <module>
     26 
     27 import torchvision
---> 28 from torchvision.models.mobilenet import ConvBNReLU
     29 from torchvision.ops import box_iou, nms
     30 

ImportError: cannot import name 'ConvBNReLU' from 'torchvision.models.mobilenet' (/usr/local/lib/python3.8/dist-packages/torchvision/models/mobilenet.py)

---------------------------------------------------------------------------
NOTE: If your import is failing due to a missing package, you can
manually install dependencies using either !pip or !apt.

To view examples of installing some common dependencies, click the
"Open Examples" button below.
---------------------------------------------------------------------------

Expected behaviour:

It should just return imported.

What did I try:

I tried to run the code in a local environment using Jupyter and my local GPU and also connected the local runtime to Google Colab by using Jupyter to use the local GPU. In both cases, the error to import persists.

PyTorch set up in the local environment confirmed that the CUDA and usage of the local GPU are enabled and available for use.

When I set the runtime to CPU, everything gets imported and works as expected, and I can train and export the model.

Therefore, I believe the issue relates to the usage of GPU and

import 'ConvBNReLU' from 'torchvision.models.mobilenet’.

I’d be grateful for any suggestions.

CodePudding user response：

Try:

from torchvision.models.mobilenetv2 import ConvBNReLU