Solved: Can't use CPU environment on exported model

(Among other issues) I’m trying to use my exported model on the CPU for my server. Currently it will throw a CUDA error, despite the model being exported to the CPU. For reference I’m following the server.py from the Render example:

async def setup_learner():
    await download_file(export_file_url, path / export_file_name)
    try:
        learn = torch.load(path/export_file_name, map_location=torch.device('cpu'))
        return learn
    except RuntimeError as e:
        if len(e.args) > 0 and 'CPU-only machine' in e.args[0]:
            print(e)
            message = "\n\nThis model was trained with an old version of fastai and will not work in a CPU environment.\n\nPlease update the fastai library in your training environment and export your model again.\n\nSee instructions for 'Returning to work' at https://course.fast.ai."
            raise RuntimeError(message)
        else:
            raise

@app.route('/analyze', methods=['POST'])
async def analyze(request):
  img_data = await request.form()
  img_bytes = await (img_data['file'].read())
  img_np = np.array(Image.open(BytesIO(img_bytes)))
  pred = learn.predict(img_np)[0]
  return JSONResponse({
      'results': str(pred)
  })

Here’s the error log:

  File "app/server.py", line 74, in analyze
    pred = learn.predict(img_np)[0]
  File "/home/muellerzr/.local/lib/python3.6/site-packages/fastai2/learner.py", line 215, in predict
    inp,preds,_,dec_preds = self.get_preds(dl=dl, with_input=True, with_decoded=True)
  File "/home/muellerzr/.local/lib/python3.6/site-packages/fastai2/learner.py", line 202, in get_preds
    self(event.begin_epoch if inner else _before_epoch)
  File "/home/muellerzr/.local/lib/python3.6/site-packages/fastai2/learner.py", line 108, in __call__
    def __call__(self, event_name): L(event_name).map(self._call_one)
  File "/home/muellerzr/.local/lib/python3.6/site-packages/fastcore/foundation.py", line 362, in map
    return self._new(map(g, self))
  File "/home/muellerzr/.local/lib/python3.6/site-packages/fastcore/foundation.py", line 315, in _new
    def _new(self, items, *args, **kwargs): return type(self)(items, *args, use_list=None, **kwargs)
  File "/home/muellerzr/.local/lib/python3.6/site-packages/fastcore/foundation.py", line 41, in __call__
    res = super().__call__(*((x,) + args), **kwargs)
  File "/home/muellerzr/.local/lib/python3.6/site-packages/fastcore/foundation.py", line 306, in __init__
    items = list(items) if use_list else _listify(items)
  File "/home/muellerzr/.local/lib/python3.6/site-packages/fastcore/foundation.py", line 242, in _listify
    if is_iter(o): return list(o)
  File "/home/muellerzr/.local/lib/python3.6/site-packages/fastcore/foundation.py", line 208, in __call__
    return self.fn(*fargs, **kwargs)
  File "/home/muellerzr/.local/lib/python3.6/site-packages/fastai2/learner.py", line 111, in _call_one
    [cb(event_name) for cb in sort_by_run(self.cbs)]
  File "/home/muellerzr/.local/lib/python3.6/site-packages/fastai2/learner.py", line 111, in <listcomp>
    [cb(event_name) for cb in sort_by_run(self.cbs)]
  File "/home/muellerzr/.local/lib/python3.6/site-packages/fastai2/callback/core.py", line 23, in __call__
    if self.run and _run: getattr(self, event_name, noop)()
  File "/home/muellerzr/.local/lib/python3.6/site-packages/fastai2/callback/core.py", line 43, in begin_fit
    self.model.to(self.dls.device)
  File "/home/muellerzr/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 425, in to
    return self._apply(convert)
  File "/home/muellerzr/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 201, in _apply
    module._apply(fn)
  File "/home/muellerzr/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 201, in _apply
    module._apply(fn)
  File "/home/muellerzr/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 223, in _apply
    param_applied = fn(param)
  File "/home/muellerzr/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 423, in convert
    return t.to(device, dtype if t.is_floating_point() else None, non_blocking)
  File "/home/muellerzr/.local/lib/python3.6/site-packages/torch/cuda/__init__.py", line 196, in _lazy_init
    _check_driver()
  File "/home/muellerzr/.local/lib/python3.6/site-packages/torch/cuda/__init__.py", line 101, in _check_driver
    http://www.nvidia.com/Download/index.aspx""")
AssertionError:
Found no NVIDIA driver on your system. Please check that you
have an NVIDIA GPU and installed a driver from
http://www.nvidia.com/Download/index.aspx

Now to make it readable, when I try to do learn.predict it still believes that I’m trying to utilize the GPU for the task and not the CPU. I’ve made sure that the learner is using the CPU, any advice? (@sgugger)

Edit: okay so in five seconds I found the solution, however I’m still keeping this up as it is a bug. The Dataloaders are not using the input device specified in load_learner. I solved it by doing learn.dls.device = 'cpu'

Also, somewhat related:

Where do I find the type-annotation for predict? As BytesIO should be added :slight_smile: (We already have it for PILImage on how to use bytes)

Ah yes, since learn.dls.device is serialized with the rest, it probably saved the cuda device. Will fix in a bit.
predict uses the type that transforms have said to it are acceptable. By default, it’s the types received during training. To add more make a type annotation on the transform you know can also accept that type, in this case PILImage.create. Look at vision.core.

1 Like

Why are you not using load_learner, that does have the cpu flag and put those dls back on the CPU for you?

1 Like

Good point, I was following the lesson 2 notebook, let me try that out! (it used torch.load)

It works, probably should have checked there first :sweat_smile: Probably a PR to the lesson2_download notebook should be done. I’ll see if I can get to it today

Yeah, this was done a while ago before load_learner was written. In general now, you should use the book draft notebooks as examples instead of the lessons, since the code there is better and up to date with all the functionality of fastai v2.

1 Like

Got it :slight_smile: Thanks!!

Just wanted to comment on this, the answer is we don’t actually have to, as we have a type for Bytes! So we can just pass in the raw byte data :slight_smile:

if isinstance(fn,bytes): fn = io.BytesIO(fn)