pytorch-torchvision.models

1.

use_cuda = torch.cuda.is_available()

vgg16 = models.vgg16(pretrained=True)

vgg16 = vgg16.cuda() if use_cuda else vgg16

img numpy

dtype = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor

img = Variable(torch.from_numpy(img)).type(dtype)

or

img = Variable(torch.from_numpy(img))

img = img.cuda()

res = vgg16(img)

res.data.cpu().numpy()