# Copyright 2019 Baidu Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import absolute_import
import numpy as np
import warnings
from perceptron.models.base import DifferentiableModel
[docs]class PyTorchModel(DifferentiableModel):
"""Creates a :class:`Model` instance from a `PyTorch` module.
Parameters
----------
model : `torch.nn.Module`
The PyTorch model that are loaded.
bounds : tuple
Tuple of lower and upper bound for the pixel values, usually
(0, 1) or (0, 255).
num_classes : int
Number of classes for which the model will output predictions.
channel_axis : int
The index of the axis that represents color channels.
device : string
A string specifying the device to do computation on.
If None, will default to "cuda:0" if torch.cuda.is_available()
or "cpu" if not.
preprocessing: 2-element tuple with floats or numpy arrays
Elementwises preprocessing of input; we first subtract the first
element of preprocessing from the input and then divide the input by
the second element.
"""
def __init__(
self,
model,
bounds,
num_classes,
channel_axis=1,
device=None,
preprocessing=(0, 1)):
# lazy import
import torch
super(PyTorchModel, self).__init__(bounds=bounds,
channel_axis=channel_axis,
preprocessing=preprocessing)
self._num_classes = num_classes
self._task = 'cls'
if device is None:
self.device = torch.device(
"cuda:0" if torch.cuda.is_available() else "cpu")
elif isinstance(device, str):
self.device = torch.device(device)
else:
self.device = device
self._model = model.to(self.device)
if model.training:
warnings.warn(
'The PyTorch model is in training model and therefore'
'might not be deterministic. Call the eval() method to'
'set it in evaluation mode if this is not intended.')
[docs] def batch_predictions(self, images):
"""Batch prediction of images."""
# lazy import
import torch
images, _ = self._process_input(images)
n = len(images)
images = torch.from_numpy(images).to(self.device)
predictions = self._model(images)
predictions = predictions.to("cpu").detach()
predictions = predictions.numpy()
assert predictions.ndim == 2
assert predictions.shape == (n, self.num_classes())
return predictions
[docs] def num_classes(self):
"""Return number of classes."""
return self._num_classes
[docs] def model_task(self):
"""Get the task that the model is used for."""
return self._task
[docs] def predictions_and_gradient(self, image, label):
"""Returns both predictions and gradients."""
# lazy import
import torch
import torch.nn as nn
input_shape = image.shape
image, dpdx = self._process_input(image)
target = np.array([label])
target = torch.from_numpy(target).long().to(self.device)
images = image[np.newaxis]
images = torch.from_numpy(images).to(self.device)
images.requires_grad_()
predictions = self._model(images)
ce = nn.CrossEntropyLoss()
loss = ce(predictions, target)
loss.backward()
grad = images.grad
predictions = predictions.to("cpu").detatch().numpy()
predictions = np.squeeze(predictions, axis=0)
assert predictions.ndim == 1
assert predictions.shape == (self.num_classes(),)
grad = grad.to("cpu").detach().numpy()
grad = np.squeeze(grad, axis=0)
grad = self._process_gradient(dpdx, grad)
assert grad.shape == input_shape
return predictions, grad
def _loss_fn(self, image, label):
# lazy import
import torch
import torch.nn as nn
image, _ = self._process_input(image)
target = np.array([label])
target = torch.from_numpy(target).long().to(self.device)
images = torch.from_numpy(image[None]).to(self.device)
predictions = self._model(images)
ce = nn.CrossEntropyLoss()
loss = ce(predictions, target)
loss = loss.to("cpu")
loss = loss.numpy()
return loss
[docs] def backward(self, gradient, image):
"""Get gradients w.r.t. the original image."""
# lazy import
import torch
assert gradient.ndim == 1
gradient = torch.from_numpy(gradient).to(self.device)
input_shape = image.shape
image, dpdx = self._process_input(image)
images = image[np.newaxis]
images = torch.from_numpy(images).to(self.device)
images.requires_grad_()
predictions = self._model(images)
predictions = predictions[0]
assert gradient.dim() == 1
assert predictions.dim() == 1
assert gradient.size() == predictions.size()
loss = torch.dot(predictions, gradient)
loss.backward()
# should be the same as predictions.backward(gradient=gradient)
grad = images.grad
grad = grad.to("cpu").detach().numpy()
grad = np.squeeze(grad, axis=0)
grad = self._process_gradient(dpdx, grad)
assert grad.shape == input_shape
return grad