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 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 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 test_stack_in_memory_single(a_tensor): iterations = 10 out = array_utils.stack_in_memory(a_tensor, iterations=iterations) assert out.shape == (10 * iterations, 3, 32, 32)