Track peak memory used on each GPU device every epoch, by :mod:`tf.contrib.memory_stats`. The peak memory comes from the ``MaxBytesInUse`` op, which is the peak memory used in recent ``session.run`` calls. See https://github.com/tensorflow/tensorflow/pull/13107.
| 235 | |
| 236 | |
| 237 | class GPUMemoryTracker(Callback): |
| 238 | """ |
| 239 | Track peak memory used on each GPU device every epoch, by :mod:`tf.contrib.memory_stats`. |
| 240 | The peak memory comes from the ``MaxBytesInUse`` op, which is the peak memory used |
| 241 | in recent ``session.run`` calls. |
| 242 | See https://github.com/tensorflow/tensorflow/pull/13107. |
| 243 | """ |
| 244 | |
| 245 | _chief_only = False |
| 246 | |
| 247 | def __init__(self, devices=(0,)): |
| 248 | """ |
| 249 | Args: |
| 250 | devices([int] or [str]): list of GPU devices to track memory on. |
| 251 | """ |
| 252 | assert isinstance(devices, (list, tuple)), devices |
| 253 | devices = ['/gpu:{}'.format(x) if isinstance(x, int) else x for x in devices] |
| 254 | self._devices = devices |
| 255 | self._disabled = False |
| 256 | |
| 257 | def _setup_graph(self): |
| 258 | try: |
| 259 | from tensorflow.contrib.memory_stats import MaxBytesInUse |
| 260 | except ImportError: |
| 261 | logger.warning("GPUMemoryTracker is not available in TF2.") |
| 262 | self._disabled = True |
| 263 | return |
| 264 | ops = [] |
| 265 | for dev in self._devices: |
| 266 | with tf.device(dev): |
| 267 | ops.append(MaxBytesInUse()) |
| 268 | self._fetches = tf.train.SessionRunArgs(fetches=ops) |
| 269 | |
| 270 | def _before_train(self): |
| 271 | if not gpu_available_in_session(): |
| 272 | self._disabled = True |
| 273 | logger.warning("GPUMemoryTracker only supports GPU!") |
| 274 | |
| 275 | def _before_run(self, _): |
| 276 | if not self._disabled and self.local_step == self.trainer.steps_per_epoch - 1: |
| 277 | return self._fetches |
| 278 | return None |
| 279 | |
| 280 | def _after_run(self, _, rv): |
| 281 | results = rv.results |
| 282 | if results is not None: |
| 283 | for mem, dev in zip(results, self._devices): |
| 284 | self.trainer.monitors.put_scalar('PeakMemory(MB)' + dev, mem / 1e6) |
| 285 | |
| 286 | |
| 287 | PeakMemoryTracker = GPUMemoryTracker |