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())
