| 290 | |
| 291 | class ImageLogger(Callback): |
| 292 | def __init__(self, batch_frequency, max_images, clamp=True, increase_log_steps=True, |
| 293 | rescale=True, disabled=False, log_on_batch_idx=False, log_first_step=False, |
| 294 | log_images_kwargs=None): |
| 295 | super().__init__() |
| 296 | self.rescale = rescale |
| 297 | self.batch_freq = batch_frequency |
| 298 | self.max_images = max_images |
| 299 | self.logger_log_images = { |
| 300 | pl.loggers.TestTubeLogger: self._testtube, |
| 301 | } |
| 302 | self.log_steps = [2 ** n for n in range(int(np.log2(self.batch_freq)) + 1)] |
| 303 | if not increase_log_steps: |
| 304 | self.log_steps = [self.batch_freq] |
| 305 | self.clamp = clamp |
| 306 | self.disabled = disabled |
| 307 | self.log_on_batch_idx = log_on_batch_idx |
| 308 | self.log_images_kwargs = log_images_kwargs if log_images_kwargs else {} |
| 309 | self.log_first_step = log_first_step |
| 310 | |
| 311 | @rank_zero_only |
| 312 | def _testtube(self, pl_module, images, batch_idx, split): |