Add option to enable/disable tensorboard (default: enable)

This commit is contained in:
allegroai 2021-04-25 10:42:50 +03:00
parent f9e555b464
commit dc6c8cfddc
2 changed files with 17 additions and 11 deletions

View File

@ -11,7 +11,6 @@ from typing import Any
import numpy as np
import six
from PIL import Image
from pathlib2 import Path
from ...debugging.log import LoggerRoot
from ..frameworks import _patched_call, WeightsFileHandler, _Empty
@ -27,15 +26,17 @@ except ImportError:
class TensorflowBinding(object):
@classmethod
def update_current_task(cls, task):
def update_current_task(cls, task, save_models, report_tensorboard):
if not task:
IsTensorboardInit.clear_tensorboard_used()
EventTrainsWriter.update_current_task(task)
PatchSummaryToEventTransformer.update_current_task(task)
PatchTensorFlowEager.update_current_task(task)
PatchKerasModelIO.update_current_task(task)
PatchTensorflowModelIO.update_current_task(task)
PatchTensorflow2ModelIO.update_current_task(task)
if report_tensorboard:
EventTrainsWriter.update_current_task(task)
PatchSummaryToEventTransformer.update_current_task(task)
PatchTensorFlowEager.update_current_task(task)
if save_models:
PatchKerasModelIO.update_current_task(task)
PatchTensorflowModelIO.update_current_task(task)
PatchTensorflow2ModelIO.update_current_task(task)
class IsTensorboardInit(object):

View File

@ -354,7 +354,7 @@ class Task(_Task):
.. code-block:: py
auto_connect_frameworks={'matplotlib': True, 'tensorflow': True, 'pytorch': True,
auto_connect_frameworks={'matplotlib': True, 'tensorflow': True, 'tensorboard': True, 'pytorch': True,
'xgboost': True, 'scikit': True, 'fastai': True, 'lightgbm': True, 'hydra': True}
:param bool auto_resource_monitoring: Automatically create machine resource monitoring plots
@ -542,9 +542,14 @@ class Task(_Task):
PatchedJoblib.update_current_task(task)
if is_auto_connect_frameworks_bool or auto_connect_frameworks.get('matplotlib', True):
PatchedMatplotlib.update_current_task(Task.__main_task)
if is_auto_connect_frameworks_bool or auto_connect_frameworks.get('tensorflow', True):
if is_auto_connect_frameworks_bool or auto_connect_frameworks.get('tensorflow', True) \
or auto_connect_frameworks.get('tensorboard', True):
PatchAbsl.update_current_task(Task.__main_task)
TensorflowBinding.update_current_task(task)
TensorflowBinding.update_current_task(task,
is_auto_connect_frameworks_bool or
auto_connect_frameworks.get('tensorflow', True),
is_auto_connect_frameworks_bool or
auto_connect_frameworks.get('tensorboard', True))
if is_auto_connect_frameworks_bool or auto_connect_frameworks.get('pytorch', True):
PatchPyTorchModelIO.update_current_task(task)
if is_auto_connect_frameworks_bool or auto_connect_frameworks.get('xgboost', True):