Skip to content

Commit

Permalink
Refactor Loggers : Move code outside train.py (ultralytics#9241)
Browse files Browse the repository at this point in the history
* update

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* update

* update

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
  • Loading branch information
3 people authored and Clay Januhowski committed Sep 8, 2022
1 parent 25a4fc5 commit cfbd364
Show file tree
Hide file tree
Showing 2 changed files with 16 additions and 6 deletions.
11 changes: 5 additions & 6 deletions train.py
Original file line number Diff line number Diff line change
Expand Up @@ -91,17 +91,16 @@ def train(hyp, opt, device, callbacks): # hyp is path/to/hyp.yaml or hyp dictio
data_dict = None
if RANK in {-1, 0}:
loggers = Loggers(save_dir, weights, opt, hyp, LOGGER) # loggers instance
if loggers.clearml:
data_dict = loggers.clearml.data_dict # None if no ClearML dataset or filled in by ClearML
if loggers.wandb:
data_dict = loggers.wandb.data_dict
if resume:
weights, epochs, hyp, batch_size = opt.weights, opt.epochs, opt.hyp, opt.batch_size

# Register actions
for k in methods(loggers):
callbacks.register_action(k, callback=getattr(loggers, k))

# Process custom dataset artifact link
data_dict = loggers.remote_dataset
if resume: # If resuming runs from remote artifact
weights, epochs, hyp, batch_size = opt.weights, opt.epochs, opt.hyp, opt.batch_size

# Config
plots = not evolve and not opt.noplots # create plots
cuda = device.type != 'cpu'
Expand Down
11 changes: 11 additions & 0 deletions utils/loggers/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -107,6 +107,17 @@ def __init__(self, save_dir=None, weights=None, opt=None, hyp=None, logger=None,
else:
self.clearml = None

@property
def remote_dataset(self):
# Get data_dict if custom dataset artifact link is provided
data_dict = None
if self.clearml:
data_dict = self.clearml.data_dict
if self.wandb:
data_dict = self.wandb.data_dict

return data_dict

def on_train_start(self):
# Callback runs on train start
pass
Expand Down

0 comments on commit cfbd364

Please sign in to comment.