def train_on_batch(self, data, target, optimizer, cuda=False, regularizer: Optional[Callable] = None): """ Train the current model on a batch using `optimizer`. Args: data (Tensor): The model input. target (Tensor): The ground truth. optimizer (optim.Optimizer): An optimizer. cuda (bool): Use CUDA or not. regularizer (Optional[Callable]): The loss regularization for training. Returns: Tensor, the loss computed from the criterion. """ if cuda: data, target = to_cuda(data), to_cuda(target) optimizer.zero_grad() output = self.model(data) loss = self.criterion(output, target) if regularizer: regularized_loss = loss + regularizer() regularized_loss.backward() else: loss.backward() optimizer.step() self._update_metrics(output, target, loss, filter='train') return loss
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 ensemble_prediction( data: torch.Tensor, model: nn.Module, weights: List[OrderedDict[str, Tensor]], cuda: bool ): """ Get the model's prediction on a batch. Args: data (Tensor): The model input. model (nn.Module): The model to use. weights (List[Dict]): List of all weights to use. 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 = to_cuda(data) res = [] for w in weights: model.load_state_dict(w) res.append(model(data)) out = _stack_preds(res) 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: 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 train_on_batch(self, data, target, optimizer, cuda=False): """ Train the current model on a batch using `optimizer`. Args: data (Tensor): the model input target (Tensor): the ground truth optimizer (optim.Optimizer): an optimizer cuda (bool): use cuda or not Returns: Tensor, the loss computed from the criterion. """ if cuda: data, target = to_cuda(data), to_cuda(target) optimizer.zero_grad() output = self.model(data) loss = self.criterion(output, target) loss.backward() optimizer.step() self._update_metrics(output, target, loss, filter='train') return loss
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()
def test_to_cuda(): t = torch.randn(3, 5) assert to_cuda(t).device.type == "cuda" m = torch.nn.Linear(3, 1) assert next(to_cuda(m).parameters()).device.type == "cuda" t_list = [t, t.clone()] assert type(t_list) is type(to_cuda(t_list)) assert all(t_.device.type == "cuda" for t_ in to_cuda(t_list)) t_dict = {"t": t, "t2": t.clone()} assert type(t_dict) is type(to_cuda(t_dict)) assert all(t_.device.type == "cuda" for t_ in to_cuda(t_dict).values()) MyTuple = namedtuple("MyTuple", ["t", "t2"]) t_named_tuple = MyTuple(t, t.clone()) assert type(t_named_tuple) is type(to_cuda(t_named_tuple)) assert all(t_.device.type == "cuda" for t_ in to_cuda(t_named_tuple)) t_tuple = (t, t.clone()) assert type(t_tuple) is type(to_cuda(t_tuple)) assert all(t_.device.type == "cuda" for t_ in to_cuda(t_tuple)) # test a type that's not explicitly defined: t_ordered_dict = OrderedDict([("t", t), ("t2", t.clone())]) assert type(t_ordered_dict) is type(to_cuda(t_ordered_dict)) assert all(t_.device.type == "cuda" for t_ in to_cuda(t_ordered_dict).values())