/mnt/d/lib/python3.7/site-packages/torch/cuda/__init__.py:52: UserWarning: CUDA initialization: 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 (Triggered internally at  /pytorch/c10/cuda/CUDAFunctions.cpp:100.)
  return torch._C._cuda_getDeviceCount() > 0

Lesson Video:

This article is also a Jupyter Notebook available to be run from the top down. There will be code snippets that you can then run in any environment.

Below are the versions of fastai, fastcore, wwf, and nbdev currently running at the time of writing this:

  • fastai: 2.1.10
  • fastcore: 1.3.13
  • wwf: 0.0.7
  • nbdev: 1.1.5

Let's now take what we had before and run inference based on a list of filenames. We'll make a quick script to get the ball rolling for how we want everything to do using nbdev again

We'll want the libraries we've used

from fastai.vision.all import *

Including our new style_transfer.py file

from wwf.style_transfer import *

Let's grab our original style image

learn = load_learner('myModel', cpu=False)

And now we can make and prepare our dataloader with a filename!

dset = Datasets('cat.jpg', tfms=[PILImage.create])
dl = dset.dataloaders(after_item=[ToTensor()], after_batch=[IntToFloatTensor(), Normalize.from_stats(*imagenet_stats)], bs=1)
t_im = dl.one_batch()[0]

And get our raw output.

with torch.no_grad():
  res = learn.model(t_im)

Let's wrap this into a function


get_learner(fn, cpu=False)

def get_learner(fn, cpu=False):
  return load_learner(fn, cpu=cpu)


make_datasets(learn, fns, bs=1)

def make_datasets(learn, fns, bs=1):
  cuda = next(learn.model.parameters()).is_cuda
  dset = Datasets(fns, tfms=[PILImage.create])
  if cuda: 
    after_batch = [IntToFloatTensor(), Normalize.from_stats(*imagenet_stats)] 
    dl = dset.dataloaders(after_item=[ToTensor()], after_batch=after_batch, bs=1)
    after_batch = [Normalize.from_stats(*imagenet_stats, cuda=False)]
    dl = dset.dataloaders(after_item=[ToTensor()], after_batch=after_batch, bs=1, device='cpu')
  return dl
from torchvision.utils import save_image

We can write a quick save_im function to save all our outputed tensors to images


save_im(imgs:list, path)

Save a ncw*h Tensor into seperate images

def save_im(imgs:list, path):
  "Save a n*c*w*h `Tensor` into seperate images"
  [save_image(im, f'{path}/{i}.png') for i, im in enumerate(imgs)]

Now let's put it all together


inference(pkl_name, fnames:list, path:Path, cpu:bool=True)

Grab inference on a model, filenames, and a path to save it to

def inference(pkl_name, fnames:list, path:Path, cpu:bool=True):
  "Grab inference on a model, filenames, and a path to save it to"
  path = path/'results'
  path.mkdir(parents=True, exist_ok=True)
  learn = get_learner(pkl_name, cpu)
  if len(fnames) > 1:
    dls = []
    for fname in fnames:
      dls.append(make_datasets(learn, fnames, 1))
    dls = [make_datasets(learn, fnames, 1)]
  res = []
  for b in dls:
    t_im = b.one_batch()[0]
    with torch.no_grad():
      out = learn.model(t_im)
  save_im(res, path)

And try it out!

fnames = ['cat.jpg'] * 5
inference('myModel', fnames, path=Path(''))

Lastly let's make a .py file again to run it off of

And we're done!