def predict_on_batch(self, data, iterations=1, cuda=False):
"""
Get the model's prediction on a batch.
Args:
data (Tensor): the model input
iterations (int): number of prediction to perform.
cuda (bool): use cuda or not
Returns:
Tensor, the loss computed from the criterion.
shape = {batch_size, nclass, n_iteration}
"""
with torch.no_grad():
if cuda:
data = data.cuda()
input_shape = data.size()
batch_size = input_shape[0]
data = torch.stack([data] * iterations)
data = data.view(batch_size * iterations, *input_shape[1:])
out = self.model(data)
out = map_on_tensor(
lambda o: o.view([iterations, batch_size, *o.size()[1:]]), out)
out = map_on_tensor(
lambda o: o.permute(1, 2, *range(3, o.ndimension()), 0), out)
return out
def mc_inference(model, data, iterations, replicate_in_memory):
if replicate_in_memory:
input_shape = data.size()
batch_size = input_shape[0]
try:
data = torch.stack([data] * iterations)
except RuntimeError as e:
raise RuntimeError(
'''CUDA ran out of memory while BaaL tried to replicate data. See the exception above.
Use `replicate_in_memory=False` in order to reduce the memory requirements.
Note that there will be some speed trade-offs''') from e
data = data.view(batch_size * iterations, *input_shape[1:])
try:
out = model(data)
except RuntimeError as e:
raise RuntimeError(
'''CUDA ran out of memory while BaaL tried to replicate data. See the exception above.
Use `replicate_in_memory=False` in order to reduce the memory requirements.
Note that there will be some speed trade-offs''') from e
out = map_on_tensor(
lambda o: o.view([iterations, batch_size, *o.size()[1:]]), out)
out = map_on_tensor(
lambda o: o.permute(1, 2, *range(3, o.ndimension()), 0), out)
else:
out = [model(data) for _ in range(iterations)]
if isinstance(out[0], Sequence):
out = [torch.stack(ts, dim=-1) for ts in zip(*out)]
else:
out = torch.stack(out, dim=-1)
return out
def test_map_on_tensor():
x = [torch.zeros([10]), torch.zeros([10])]
assert np.allclose(map_on_tensor(lambda xi: (xi + 1).numpy(), x),
[np.ones([10]), np.ones([10])])
x = torch.zeros([10])
assert np.allclose(
map_on_tensor(lambda xi: xi + 1, x).numpy(), np.ones([10]))
def predict_on_batch(self, data, iterations=1, cuda=False):
"""
Get the model's prediction on a batch.
Args:
data (Tensor): The model input.
iterations (int): Number of prediction to perform.
cuda (bool): Use CUDA or not.
Returns:
Tensor, the loss computed from the criterion.
shape = {batch_size, nclass, n_iteration}.
Raises:
Raises RuntimeError if CUDA rans out of memory during data replication.
"""
with torch.no_grad():
if cuda:
data = to_cuda(data)
if self.replicate_in_memory:
data = map_on_tensor(lambda d: stack_in_memory(d, iterations), data)
try:
out = self.model(data)
except RuntimeError as e:
raise RuntimeError(
'''CUDA ran out of memory while BaaL tried to replicate data. See the exception above.
Use `replicate_in_memory=False` in order to reduce the memory requirements.
Note that there will be some speed trade-offs''') from e
out = map_on_tensor(lambda o: o.view([iterations, -1, *o.size()[1:]]), out)
out = map_on_tensor(lambda o: o.permute(1, 2, *range(3, o.ndimension()), 0), out)
else:
out = [self.model(data) for _ in range(iterations)]
out = _stack_preds(out)
return out
def predict_on_batch(self, data, iterations=1, cuda=False):
"""
Get the model's prediction on a batch.
Args:
data (Tensor): the model input
iterations (int): number of prediction to perform.
cuda (bool): use cuda or not
Returns:
Tensor, the loss computed from the criterion.
shape = {batch_size, nclass, n_iteration}
Raises:
raises RuntimeError if CUDA rans out of memory during data replication.
"""
with torch.no_grad():
if cuda:
data = to_cuda(data)
if self.replicate_in_memory:
input_shape = data.size()
batch_size = input_shape[0]
try:
data = torch.stack([data] * iterations)
except RuntimeError as e:
raise RuntimeError(
'''CUDA ran out of memory while BaaL tried to replicate data. See the exception above.
Use `replicate_in_memory=False` in order to reduce the memory requirements.
Note that there will be some speed trade-offs''') from e
data = data.view(batch_size * iterations, *input_shape[1:])
try:
out = self.model(data)
except RuntimeError as e:
raise RuntimeError(
'''CUDA ran out of memory while BaaL tried to replicate data. See the exception above.
Use `replicate_in_memory=False` in order to reduce the memory requirements.
Note that there will be some speed trade-offs''') from e
out = map_on_tensor(
lambda o: o.view([iterations, batch_size, *o.size()[1:]]),
out)
out = map_on_tensor(
lambda o: o.permute(1, 2, *range(3, o.ndimension()), 0),
out)
else:
out = [self.model(data) for _ in range(iterations)]
if isinstance(out[0], Sequence):
out = [torch.stack(ts, dim=-1) for ts in zip(*out)]
else:
out = torch.stack(out, dim=-1)
return out
def test_on_batch(
self,
data: torch.Tensor,
target: torch.Tensor,
cuda: bool = False,
average_predictions: int = 1,
):
"""
Test the current model on a batch.
Args:
data (Tensor): the model input
target (Tensor): the ground truth
cuda (bool): use cuda or not
average_predictions (int): The number of predictions to average to
compute the test loss.
Returns:
Tensor, the loss computed from the criterion.
"""
with torch.no_grad():
if cuda:
data, target = to_cuda(data), to_cuda(target)
if average_predictions == 1:
preds = self.model(data)
loss = self.criterion(preds, target)
elif average_predictions > 1:
preds = map_on_tensor(
lambda p: p.mean(-1),
self.predict_on_batch(data,
iterations=average_predictions,
cuda=cuda))
loss = self.criterion(preds, target)
self._update_metrics(preds, target, loss, 'test')
return loss
def predict_on_dataset_generator(self, dataloader=None, *args, **kwargs):
"""Predict on the pool loader.
Args:
dataloader (Optional[DataLoader]): If provided, will predict on this dataloader.
Otherwise, uses model.pool_loader().
Returns:
Numpy arrays with all the predictions.
"""
model = self.get_model()
model.eval()
if self.on_gpu:
model.cuda(self.root_gpu)
dataloader = dataloader or model.pool_loader()
if len(dataloader) == 0:
return None
log.info("Start Predict", dataset=len(dataloader))
for idx, batch in enumerate(tqdm(dataloader, total=len(dataloader), file=sys.stdout)):
if self.on_gpu:
batch = to_cuda(batch)
pred = self.model.predict_step(batch, idx)
yield map_on_tensor(lambda x: x.detach().cpu().numpy(), pred)
# teardown, TODO customize this later?
model.cpu()
def predict_on_dataset_generator(self,
model=None,
dataloader: Optional[DataLoader] = None,
*args,
**kwargs):
"""Predict on the pool loader.
Args:
model: Model to be used in prediction. If None, will get the Trainer's model.
dataloader (Optional[DataLoader]): If provided, will predict on this dataloader.
Otherwise, uses model.pool_dataloader().
Returns:
Numpy arrays with all the predictions.
"""
model = model or self.lightning_module
model.eval()
if isinstance(self.accelerator, GPUAccelerator):
model.cuda(self.accelerator.root_device)
dataloader = dataloader or model.pool_dataloader()
if len(dataloader) == 0:
return None
log.info("Start Predict", dataset=len(dataloader))
for idx, batch in enumerate(
tqdm(dataloader, total=len(dataloader), file=sys.stdout)):
if isinstance(self.accelerator, GPUAccelerator):
batch = to_cuda(batch)
pred = model.predict_step(batch, idx)
yield map_on_tensor(lambda x: x.detach().cpu().numpy(), pred)
# teardown, TODO customize this later?
model.cpu()
def test_stack_in_memory_multi(a_tensor):
iterations = 10
t = [a_tensor, a_tensor]
out = map_on_tensor(
lambda ti: array_utils.stack_in_memory(ti, iterations=iterations), t)
assert out[0].shape == (10 * iterations, 3, 32, 32)
assert out[1].shape == (10 * iterations, 3, 32, 32)
def test_map_on_dict():
x = {'key': torch.zeros([10]),
'door': torch.zeros([10])}
result = map_on_tensor(lambda xi: (xi + 1).numpy(), x)
assert isinstance(result, dict)
for _, v in result.items():
assert np.allclose(v, np.ones([10]))
def predict_on_dataset_generator(
self,
dataset: Dataset,
batch_size: int,
iterations: int,
use_cuda: bool,
workers: int = 4,
collate_fn: Optional[Callable] = None,
half=False,
verbose=True,
):
"""
Use the model to predict on a dataset `iterations` time.
Args:
dataset (Dataset): Dataset to predict on.
batch_size (int): Batch size to use during prediction.
iterations (int): Number of iterations per sample.
use_cuda (bool): Use CUDA or not.
workers (int): Number of workers to use.
collate_fn (Optional[Callable]): The collate function to use.
half (bool): If True use half precision.
verbose (bool): If True use tqdm to display progress
Notes:
The "batch" is made of `batch_size` * `iterations` samples.
Returns:
Generators [batch_size, n_classes, ..., n_iterations].
"""
self.eval()
if len(dataset) == 0:
return None
log.info("Start Predict", dataset=len(dataset))
collate_fn = collate_fn or default_collate
loader = DataLoader(dataset, batch_size, False, num_workers=workers, collate_fn=collate_fn)
if verbose:
loader = tqdm(loader, total=len(loader), file=sys.stdout)
for idx, (data, _) in enumerate(loader):
pred = self.predict_on_batch(data, iterations, use_cuda)
pred = map_on_tensor(lambda x: x.detach(), pred)
if half:
pred = map_on_tensor(lambda x: x.half(), pred)
yield map_on_tensor(lambda x: x.cpu().numpy(), pred)
def predict_on_dataset_generator(
self,
dataset,
iterations: int = 1,
half: bool = False,
ignore_keys: Optional[List[str]] = None,
):
"""
Use the model to predict on a dataset `iterations` time.
Args:
dataset (Dataset): Dataset to predict on.
iterations (int): Number of iterations per sample.
half (bool): If True use half precision.
ignore_keys (Optional[List[str]]): A list of keys in the output of your model
(if it is a dictionary) that should be ignored when gathering predictions.
Notes:
The "batch" is made of `batch_size` * `iterations` samples.
Returns:
Generators [batch_size, n_classes, ..., n_iterations].
"""
dataloader = self.get_eval_dataloader(dataset)
log.info("Start Predict", dataset=len(dataset))
model = self.model
model.eval()
for step, inputs in enumerate(tqdm(dataloader)):
inputs = map_on_tensor(
lambda element: map_on_tensor(
lambda d: stack_in_memory(d, iterations), element),
inputs,
)
_, out, _ = self.prediction_step(model,
inputs,
prediction_loss_only=False,
ignore_keys=ignore_keys)
out = map_on_tensor(
lambda o: o.view([iterations, -1, *o.size()[1:]]), out)
out = map_on_tensor(
lambda o: o.permute(1, 2, *range(3, o.ndimension()), 0), out)
out = map_on_tensor(lambda x: x.detach(), out)
if half:
out = map_on_tensor(lambda x: x.half(), out)
yield map_on_tensor(lambda x: x.cpu().numpy(), out)
def predict_on_dataset_generator(self, *args, **kwargs):
model = self.get_model()
model.eval()
if self.on_gpu:
model.cuda(self.root_gpu)
dataloader = self.model.pool_loader()
if len(dataloader) == 0:
return None
log.info("Start Predict", dataset=len(dataloader))
for idx, (data, _) in enumerate(
tqdm(dataloader, total=len(dataloader), file=sys.stdout)):
if self.on_gpu:
data = to_cuda(data)
pred = self.model.predict_step(data, idx)
yield map_on_tensor(lambda x: x.detach().cpu().numpy(), pred)
# teardown, TODO customize this later?
model.cpu()
