def test_attach_model_callbacks_override_info(caplog): """Test that the logs contain the info about overriding callbacks returned by configure_callbacks.""" model = LightningModule() model.configure_callbacks = lambda: [ LearningRateMonitor(), EarlyStopping(monitor="foo") ] trainer = Trainer(enable_checkpointing=False, callbacks=[ EarlyStopping(monitor="foo"), LearningRateMonitor(), TQDMProgressBar() ]) trainer.model = model cb_connector = CallbackConnector(trainer) with caplog.at_level(logging.INFO): cb_connector._attach_model_callbacks() assert "existing callbacks passed to Trainer: EarlyStopping, LearningRateMonitor" in caplog.text
def on_epoch_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule): if trainer.current_epoch % self.snapshot_interval == 0: train_dataloader = pl_module.train_dataloader() for x, y, global_example_indexes in tqdm( train_dataloader, desc="Snapshotting: ", file=sys.stdout, position=0, leave=True, ): x.requires_grad = True logits = pl_module.forward(x) loss = F.nll_loss(logits, y) loss.backward() for global_index, grad, class_label in zip( global_example_indexes.tolist(), x.grad, y): self._x_idx_to_grads.setdefault(global_index, []).append(grad) self._x_idx_to_y[global_index] = class_label
def on_train_batch_end( self, trainer: Trainer, pl_module: LightningModule, outputs: Sequence, batch: Sequence, batch_idx: int, dataloader_idx: int, ) -> None: # show images only every 20 batches if (batch_idx + 1) % self.logging_batch_interval != 0: return # pick the last batch and logits x, y = batch try: logits = pl_module.last_logits except AttributeError as err: m = """please track the last_logits in the training_step like so: def training_step(...): self.last_logits = your_logits """ raise AttributeError(m) from err # only check when it has opinions (ie: the logit > 5) if logits.max() > self.min_logit_value: # pick the top two confused probs (values, idxs) = torch.topk(logits, k=2, dim=1) # care about only the ones that are at most eps close to each other eps = self.max_logit_difference mask = (values[:, 0] - values[:, 1]).abs() < eps if mask.sum() > 0: # pull out the ones we care about confusing_x = x[mask, ...] confusing_y = y[mask] mask_idxs = idxs[mask] pl_module.eval() self._plot(confusing_x, confusing_y, trainer, pl_module, mask_idxs) pl_module.train()
def _train_log(self, trainer: Trainer, pl_module: LightningModule): if self.training_config.evaluate_metrics: self.train_combined_report.metrics = pl_module.metrics( self.train_combined_report, self.train_combined_report) pl_module.train_meter.update_from_report(self.train_combined_report) extra = {} if "cuda" in str(trainer.model.device): extra["max mem"] = torch.cuda.max_memory_allocated() / 1024 extra["max mem"] //= 1024 if self.training_config.experiment_name: extra["experiment"] = self.training_config.experiment_name optimizer = self.get_optimizer(trainer) num_updates = self._get_num_updates_for_logging(trainer) current_iteration = self._get_iterations_for_logging(trainer) extra.update({ "epoch": self._get_current_epoch_for_logging(trainer), "iterations": current_iteration, "num_updates": num_updates, "max_updates": trainer.max_steps, "lr": "{:.5f}".format(optimizer.param_groups[0]["lr"]).rstrip("0"), "ups": "{:.2f}".format(self.trainer_config.log_every_n_steps / self.train_timer.unix_time_since_start()), "time": self.train_timer.get_time_since_start(), "time_since_start": self.total_timer.get_time_since_start(), "eta": calculate_time_left( max_updates=trainer.max_steps, num_updates=num_updates, timer=self.train_timer, num_snapshot_iterations=self.snapshot_iterations, log_interval=self.trainer_config.log_every_n_steps, eval_interval=self.trainer_config.val_check_interval, ), }) self.train_timer.reset() summarize_report( current_iteration=current_iteration, num_updates=num_updates, max_updates=trainer.max_steps, meter=pl_module.train_meter, extra=extra, tb_writer=self.lightning_trainer.tb_writer, )
def on_validation_batch_end( self, trainer: Trainer, pl_module: LightningModule, outputs: Sequence, batch: Sequence, batch_idx: int, dataloader_idx: int, ) -> None: x, y = self.to_device(batch, pl_module.device) with torch.no_grad(): representations = self.get_representations(pl_module, x) representations = representations.detach() # forward pass mlp_preds = pl_module.non_linear_evaluator( representations) # type: ignore[operator] mlp_loss = F.cross_entropy(mlp_preds, y) # log metrics val_acc = accuracy(mlp_preds, y) pl_module.log('online_val_acc', val_acc, on_step=False, on_epoch=True, sync_dist=True) pl_module.log('online_val_loss', mlp_loss, on_step=False, on_epoch=True, sync_dist=True) self.confusion_matrix(mlp_preds, y) if self.time_to_sample: N, C, H, W = batch[0][2].shape num = min(N, 16) self.images = batch[0][2][0:num] self.mlp_preds = torch.argmax(mlp_preds[0:num], dim=1) self.labels = y[0:num] self.time_to_sample = False
def make_logging(self, pl_module: pl.LightningModule, func_name, func_idx, on_steps=[], on_epochs=[], prob_bars=[]): self.funcs_called_count[func_name] += 1 for idx, (on_step, on_epoch, prog_bar) in enumerate(list(itertools.product(*[on_steps, on_epochs, prob_bars]))): # run logging custom_func_name = f"{func_idx}_{idx}_{func_name}" pl_module.log(custom_func_name, self.count * func_idx, on_step=on_step, on_epoch=on_epoch, prog_bar=prog_bar) # catch information for verification # on on_train_start is outside the main loop. Won't be called if func_name == "on_train_start": self.callback_funcs_called[func_name].append([self.count * func_idx]) # Saved only values from second epoch, so we can compute its mean or latest. if pl_module.trainer.current_epoch == 1: self.callback_funcs_called[func_name].append([self.count * func_idx]) forked = on_step and on_epoch self.funcs_attr[custom_func_name] = { "on_step": on_step, "on_epoch": on_epoch, "prog_bar": prog_bar, "forked": forked, "func_name": func_name} if on_step and on_epoch: self.funcs_attr[f"{custom_func_name}_step"] = { "on_step": True, "on_epoch": False, "prog_bar": prog_bar, "forked": False, "func_name": func_name} self.funcs_attr[f"{custom_func_name}_epoch"] = { "on_step": False, "on_epoch": True, "prog_bar": prog_bar, "forked": False, "func_name": func_name}
def on_test_start(self, trainer: Trainer, pl_module: LightningModule) -> None: """Make sure we have a quantized version. This handles the edge case where a user does .test() without .fit() first. """ if hasattr(pl_module, "_quantized"): return pl_module._quantized = self.convert( pl_module._prepared, self.qconfig_dicts.keys(), attrs=self.preserved_attrs ) self.quantized = pl_module._quantized
def on_pretrain_routine_start(self, trainer: Trainer, pl_module: LightningModule) -> None: from pl_bolts.models.self_supervised.evaluator import SSLEvaluator pl_module.non_linear_evaluator = SSLEvaluator( n_input=self.z_dim, n_classes=self.num_classes, p=self.drop_p, n_hidden=self.hidden_dim, ).to(pl_module.device) self.optimizer = torch.optim.Adam(pl_module.non_linear_evaluator.parameters(), lr=1e-4)
def fit( self, model: pl.LightningModule, train_dataloader: Optional[DataLoader] = None, val_dataloaders: Optional[DataLoader] = None, ): self.model = model # try to get hparams from module. if self.tune_hparams is None: if hasattr(model, 'get_pbt_hparams'): self.tune_hparams = model.get_pbt_hparams() mp = _mp.get_context('forkserver') global_epoch = mp.Value('i', 0) # initialize population tasks population_tasks = mp.Queue(maxsize=self.population_size) for i in range(self.population_size): population_tasks.put({ 'id': i, self.pbt_monitor: 0, 'checkpoint_path': None, }) logger_info = dict(name=self.trainer.logger.name, version=self.trainer.logger.version, save_dir=self.trainer.logger.save_dir) workers = [ Worker( pl_trainer=copy.deepcopy(self.trainer), model=copy.deepcopy(model), population_tasks=population_tasks, tune_hparams=self.tune_hparams, process_position=i, global_epoch=global_epoch, max_epoch=10, full_parallel=self.num_workers == self.population_size, logger_info=logger_info, dataloaders=dict( train_dataloader=copy.deepcopy(train_dataloader), val_dataloaders=copy.deepcopy(val_dataloaders), ), ) for i in range(self.num_workers) ] [w.start() for w in workers] [w.join() for w in workers] task = [] while not population_tasks.empty(): task.append(population_tasks.get())
def on_validation_batch_end( self, trainer: Trainer, pl_module: LightningModule, outputs: Sequence, batch: Sequence, batch_idx: int, dataloader_idx: int, ) -> None: val_acc, mlp_loss = self.shared_step(pl_module, batch) pl_module.log("online_val_acc", val_acc, on_step=False, on_epoch=True, sync_dist=True) pl_module.log("online_val_loss", mlp_loss, on_step=False, on_epoch=True, sync_dist=True)
def test_forward(self, model: pl.LightningModule, data: torch.Tensor, training: bool): r"""Calls ``model.forward()`` and tests that the output is not ``None``. Because of the size of some models, this test is only run when a GPU is available. """ if isinstance(model, (LightningDistributedDataParallel, LightningDistributedModule)): pytest.skip() if torch.cuda.is_available(): model = model.cuda() # type: ignore data = data.cuda() if training: model.train() else: model.eval() _ = model(data) assert _ is not None
def test_invalid_weights_summmary(): """Test that invalid value for weights_summary raises an error.""" model = LightningModule() with pytest.raises( MisconfigurationException, match="`weights_summary` can be None, .* got temp" ), pytest.deprecated_call(match="weights_summary=temp)` is deprecated"): Trainer(weights_summary="temp") with pytest.raises(ValueError, match="max_depth` can be .* got temp"): ModelSummary(model, max_depth="temp")
def log_hyperparameters( config: DictConfig, model: pl.LightningModule, datamodule: pl.LightningDataModule, trainer: pl.Trainer, callbacks: List[pl.Callback], logger: List[pl.loggers.LightningLoggerBase], ) -> None: """Controls which config parts are saved by Lightning loggers. Additionaly saves: - number of model parameters """ if not trainer.logger: return hparams = {} # choose which parts of hydra config will be saved to loggers hparams["model"] = config["model"] # save number of model parameters hparams["model/params/total"] = sum(p.numel() for p in model.parameters()) hparams["model/params/trainable"] = sum(p.numel() for p in model.parameters() if p.requires_grad) hparams["model/params/non_trainable"] = sum(p.numel() for p in model.parameters() if not p.requires_grad) hparams["datamodule"] = config["datamodule"] hparams["trainer"] = config["trainer"] if "seed" in config: hparams["seed"] = config["seed"] if "callbacks" in config: hparams["callbacks"] = config["callbacks"] # send hparams to all loggers trainer.logger.log_hyperparams(hparams)
def log_hyperparameters( config: DictConfig, model: pl.LightningModule, datamodule: pl.LightningDataModule, trainer: pl.Trainer, callbacks: List[pl.Callback], logger: List[pl.loggers.LightningLoggerBase], ) -> None: """This method controls which parameters from Hydra config are saved by Lightning loggers. Additionaly saves: - number of trainable model parameters """ hparams = {} # choose which parts of hydra config will be saved to loggers hparams["trainer"] = config["trainer"] hparams["model"] = config["model"] hparams["datamodule"] = config["datamodule"] if "seed" in config: hparams["seed"] = config["seed"] if "callbacks" in config: hparams["callbacks"] = config["callbacks"] # save number of model parameters hparams["model/params_total"] = sum(p.numel() for p in model.parameters()) hparams["model/params_trainable"] = sum( p.numel() for p in model.parameters() if p.requires_grad ) hparams["model/params_not_trainable"] = sum( p.numel() for p in model.parameters() if not p.requires_grad ) # send hparams to all loggers trainer.logger.log_hyperparams(hparams) # disable logging any more hyperparameters for all loggers # this is just a trick to prevent trainer from logging hparams of model, # since we already did that above trainer.logger.log_hyperparams = empty
def common_epoch_end(self, step_outputs, prefix="train/", exclude_keys={"pred", "target"}): keys = list(step_outputs[0].keys()) mean_step_metrics = { k: torch.mean(torch.stack([x[k] for x in step_outputs])) for k in keys if k not in exclude_keys } preds, targets = zip(*[(s["pred"], s["target"]) for s in step_outputs]) preds = cat_steps(preds) targets = cat_steps(targets) epoch_metrics = prefix_keys( prefix, self.collect_epoch_metrics(preds, targets, prefix.replace("/", "")), ) epoch_metrics, epoch_figures = sort_out_figures(epoch_metrics) all_metrics = {**mean_step_metrics, **epoch_metrics} LightningModule.log_dict(self, all_metrics, sync_dist=self._sync_dist) log_figures(self, epoch_figures)
def setup(self, trainer: Trainer, pl_module: LightningModule, stage: str) -> None: """Override the model with a quantized-aware version on setup. This is the earliest place we can override this model which allows for appropriate behavior when restoring from checkpoints, as well as connecting to accelerators, etc. The model is only prepared once. """ # Only prepare the model once. if hasattr(pl_module, "_prepared"): return with mode(pl_module, training=True) as train: pl_module._prepared = self.prepare( _deepcopy(train), configs=self.qconfig_dicts, attrs=self.preserved_attrs, ) pl_module.forward = MethodType(_quantized_forward, pl_module) self.prepared = pl_module._prepared
def on_train_epoch_end(self, trainer : pl.Trainer, pl_module : pl.LightningModule, outputs): if pl_module.current_epoch % self.save_step == 0: ckpt = { 'net_model' : pl_module.model.state_dict(), 'ema_model' : pl_module.ema_model.state_dict(), 'optim' : pl_module.optimizers().state_dict(), 'epoch' : pl_module.current_epoch, #'x_T' : self.x_T } Path(os.path.join(self.save_dir)).mkdir(parents=True, exist_ok=True) path = os.path.join(self.save_dir, 'ckpt.pt') torch.save(ckpt, path)
def on_train_start(self, trainer: Trainer, pl_module: LightningModule): super().on_train_start(trainer, pl_module) submodule_dict = dict(pl_module.named_modules()) self._hook_handles = [] for name in self._get_submodule_names(pl_module): if name not in submodule_dict: rank_zero_warn( f"{name} is not a valid identifier for a submodule in {pl_module.__class__.__name__}," " skipping this key.") continue handle = self._register_hook(name, submodule_dict[name]) self._hook_handles.append(handle)
def lr_find(module: pl.LightningModule, gpu_id: typing.Union[torch.device, int] = None, init_value: float = 1e-8, final_value: float = 10., beta: float = 0.98, max_steps: int = None) -> (typing.List[float], typing.List[float]): with tempfile.TemporaryDirectory() as tmpdir: save_path = pathlib.Path(tmpdir) / 'model.pth' torch.save(module.state_dict(), save_path) train_dataloader = module.train_dataloader() if max_steps is None: num = len(train_dataloader) - 1 else: num = min(len(train_dataloader) - 1, max_steps) mult = (final_value / init_value)**(1 / num) lr = init_value avg_loss = 0. best_loss = 0. losses = [] lrs = [] optimizers = initialize_optimizers(module, lr) if gpu_id is not None: module = module.to(gpu_id) for batch_num, batch in enumerate(tqdm(train_dataloader, total=num), start=1): if gpu_id is not None: batch = transfer_batch_to_gpu(batch, gpu_id) loss = module.training_step(batch, batch_num)['loss'] # Compute the smoothed loss avg_loss = beta * avg_loss + (1 - beta) * loss.item() smoothed_loss = avg_loss / (1 - beta**batch_num) # Stop if the loss is exploding if batch_num > 1 and smoothed_loss > 4 * best_loss: break if lr >= final_value: break if smoothed_loss < best_loss or batch_num == 1: best_loss = smoothed_loss losses.append(smoothed_loss) lrs.append(lr) loss.backward() optimizers = step_optimizers(optimizers) optimizers = zero_grad_optimizers(optimizers) # Update the lr for the next step lr *= mult optimizers = set_optimizer_lr(optimizers, lr) module.load_state_dict(torch.load(save_path)) return lrs, losses
def test_checkpoint_callbacks_are_last(tmpdir): """Test that checkpoint callbacks always get moved to the end of the list, with preserved order.""" checkpoint1 = ModelCheckpoint(tmpdir) checkpoint2 = ModelCheckpoint(tmpdir) model_summary = ModelSummary() early_stopping = EarlyStopping() lr_monitor = LearningRateMonitor() progress_bar = ProgressBar() # no model reference trainer = Trainer(callbacks=[checkpoint1, progress_bar, lr_monitor, model_summary, checkpoint2]) cb_connector = CallbackConnector(trainer) cb_connector._attach_model_callbacks() assert trainer.callbacks == [progress_bar, lr_monitor, model_summary, checkpoint1, checkpoint2] # no model callbacks model = LightningModule() model.configure_callbacks = lambda: [] trainer.model = model cb_connector._attach_model_callbacks() assert trainer.callbacks == [progress_bar, lr_monitor, model_summary, checkpoint1, checkpoint2] # with model-specific callbacks that substitute ones in Trainer model = LightningModule() model.configure_callbacks = lambda: [checkpoint1, early_stopping, model_summary, checkpoint2] trainer = Trainer(callbacks=[progress_bar, lr_monitor, ModelCheckpoint(tmpdir)]) trainer.model = model cb_connector = CallbackConnector(trainer) cb_connector._attach_model_callbacks() assert trainer.callbacks == [progress_bar, lr_monitor, early_stopping, model_summary, checkpoint1, checkpoint2]
def on_validation_batch_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule, outputs: Any, batch: BatchType, batch_idx: int, dataloader_idx: int) -> None: # type: ignore """ Get and log validation metrics. """ ids_linear_head = tuple( batch[SSLDataModuleType.LINEAR_HEAD][0].tolist()) if ids_linear_head not in self.visited_ids: self.visited_ids.add(ids_linear_head) loss = self.shared_step(batch, pl_module, is_training=False) pl_module.log('ssl_online_evaluator/val/loss', loss, on_step=False, on_epoch=True, sync_dist=False) for metric in self.val_metrics: pl_module.log(f"ssl_online_evaluator/val/{metric.name}", metric, on_epoch=True, on_step=False) # type: ignore
def override_unsupported_nud(lm: pl.LightningModule, context: PyTorchTrialContext) -> None: writer = pytorch.TorchWriter() def lm_print(*args: Any, **kwargs: Any) -> None: if context.distributed.get_rank() == 0: print(*args, **kwargs) def lm_log_dict(a_dict: Dict, *args: Any, **kwargs: Any) -> None: if len(args) != 0 or len(kwargs) != 0: raise InvalidModelException( f"unsupported arguments to LightningModule.log {args} {kwargs}" ) for metric, value in a_dict.items(): if type(value) == int or type(value) == float: writer.add_scalar(metric, value, context.current_train_batch()) def lm_log(name: str, value: Any, *args: Any, **kwargs: Any) -> None: lm_log_dict({name: value}, *args, **kwargs) lm.print = lm_print # type: ignore lm.log = lm_log # type: ignore lm.log_dict = lm_log_dict # type: ignore
def on_train_epoch_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None: if not _TORCHVISION_AVAILABLE: return images, _ = next( iter( DataLoader(trainer.datamodule.mnist_val, batch_size=self.num_samples))) images_flattened = images.view(images.size(0), -1) # generate images with torch.no_grad(): pl_module.eval() images_generated = pl_module(images_flattened.to(pl_module.device)) pl_module.train() if trainer.current_epoch == 0: save_image(self._to_grid(images), f"grid_ori_{trainer.current_epoch}.png") save_image(self._to_grid(images_generated.reshape(images.shape)), f"grid_generated_{trainer.current_epoch}.png")
def on_train_batch_end( self, trainer: Trainer, pl_module: LightningModule, outputs: Sequence, batch: Sequence, batch_idx: int, dataloader_idx: int, ) -> None: if trainer.current_epoch >= self.start_epoch: x, y = self.to_device(batch, pl_module.device) with torch.no_grad(): representations = self.get_representations(pl_module, x) representations = representations.detach() # forward pass mlp_preds = pl_module.non_linear_evaluator( representations) # type: ignore[operator] weight = (None if self.loss_weight is None else self.loss_weight.to(pl_module.device).float()) mlp_loss = F.cross_entropy(mlp_preds, y, weight=weight) # update finetune weights mlp_loss.backward() self.optimizer.step() self.optimizer.zero_grad() # log metrics train_acc = accuracy(mlp_preds, y) pl_module.log("online_train_acc", train_acc, on_step=True, on_epoch=False) pl_module.log("online_train_loss", mlp_loss, on_step=True, on_epoch=False)
def on_test_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule): logger.info("***** Test results *****") if pl_module.is_logger(): metrics = trainer.callback_metrics # Log and save results to file output_test_results_file = os.path.join(pl_module.hparams.output_dir, "test_results.txt") with open(output_test_results_file, "w") as writer: for key in sorted(metrics): if key not in ["log", "progress_bar"]: logger.info("{} = {}\n".format(key, str(metrics[key]))) writer.write("{} = {}\n".format(key, str(metrics[key])))
def create_figure( self, task: pl.LightningModule, x: Tensor, year: np.array = None ) -> plt.Figure: """ Creates a matplotlib figure of generated album covers. Args: task: The CoverGANTask that contains the generator that should be used to generate the images from the latent vectors x: Latent vectors year: Standardized release year of the artificial album """ if year is not None: year = Tensor(year).to(task.device) idx = 1 figsize = (np.array(self.target_size) * [10, 1]).astype(int) / 300 fig = plt.figure(figsize=figsize, dpi=300) with torch.no_grad(): if isinstance(task.generator, ProGANGenerator): output = task.generator( x, year=year, block=task.block, alpha=task.alpha ) else: output = task.generator(x) generated_images = output.detach().cpu().numpy() generated_images = np.moveaxis(generated_images, 1, -1) for img in generated_images: if img.shape[-1] == 1: img = np.tile(img, (1, 1, 3)) img = array_to_img(img, scale=True) img = img.resize(size=self.target_size) plt.subplot(1, 10, idx) plt.axis("off") plt.imshow(img) idx += 1 plt.subplots_adjust( left=0, bottom=0, right=1, top=1, wspace=0, hspace=0.1 ) return fig
def _deepcopy(pl_module: LightningModule) -> LightningModule: """Copy a LightningModule. Some properties need to be ignored. """ # Remove _result before call to deepcopy since it store non-leaf Tensors. # If not removed, you'll see this error on deepcopy() attempts: P150283141. if hasattr(pl_module, "_results"): result = pl_module._results delattr(pl_module, "_results") copy = deepcopy(pl_module) # Set back. pl_module._results = result else: copy = deepcopy(pl_module) return copy
def custom_lr_finder(runner: LightningModule, omegaConf: DictConfig) -> LightningModule: """ LR finder The prepare_data/setup does not work well with the lr finder To handle the situation we manually search for it before the training The situation is being handled: https://github.com/PyTorchLightning/pytorch-lightning/issues/2485 """ del omegaConf['trainer']['auto_lr_find'] tmp_trainer = pl.Trainer(**omegaConf['trainer']) runner.prepare_data() runner.setup('lr_finder') lr_finder = tmp_trainer.lr_find(runner) # fig = lr_finder.plot(suggest=True) new_lr = lr_finder.suggestion() omegaConf['runner']['lr'] = new_lr runner = make_runner(omegaConf['runner']) if omegaConf['runner'].get('verbose', False) is True: print('Learning rate found: {}'.format(new_lr)) return runner
def wrap_training_step(wrapped, instance: LightningModule, args, kwargs): """ Wraps the training step of the LightningModule. Parameters ---------- wrapped: The wrapped function. instance: The LightningModule instance. args: The arguments passed to the wrapped function. kwargs: The keyword arguments passed to the wrapped function. Returns ------- The return value of the wrapped function. """ output_dict = wrapped(*args, **kwargs) if isinstance(output_dict, dict) and output_dict is not None and "log" in output_dict: log_dict = output_dict.pop("log") instance.log_dict(log_dict, on_step=True) return output_dict
def test_validation_step(self, model: pl.LightningModule): r"""Runs a validation step based on the data returned from ``model.val_dataloader()``. Tests that the dictionary returned from ``validation_step()`` are as required by PyTorch Lightning. Because of the size of some models, this test is only run when a GPU is available. """ if isinstance(model, LightningDistributedDataParallel): pytest.skip() check_overriden(model, "val_dataloader") check_overriden(model, "validation_step") dl = model.val_dataloader() # TODO this can't handle multiple optimizers batch = next(iter(dl)) if torch.cuda.is_available(): batch = [x.cuda() for x in batch] model = model.cuda() model.eval() output = model.validation_step(batch, 0) assert isinstance(output, dict) if "loss" in output.keys(): assert isinstance(output["loss"], Tensor) assert output["loss"].shape == (1, ) if "log" in output.keys(): assert isinstance(output["log"], dict) if "progress_bar" in output.keys(): assert isinstance(output["progress_bar"], dict) return batch, output