henry000 commited on
Commit
9c42a35
·
1 Parent(s): a0cddf7

🎨 [Update] progress bar, update epoch each batch

Browse files
Files changed (1) hide show
  1. yolo/utils/logging_utils.py +5 -3
yolo/utils/logging_utils.py CHANGED
@@ -22,6 +22,7 @@ from rich.console import Console
22
  from rich.progress import BarColumn, Progress, TextColumn, TimeRemainingColumn
23
  from rich.table import Table
24
  from torch import Tensor
 
25
 
26
  from yolo.config.config import Config, GeneralConfig, YOLOLayer
27
 
@@ -40,7 +41,7 @@ class ProgressTracker:
40
  self.progress = Progress(
41
  TextColumn("[progress.description]{task.description}"),
42
  BarColumn(bar_width=None),
43
- TextColumn("{task.completed}/{task.total}"),
44
  TimeRemainingColumn(),
45
  )
46
  self.use_wandb = use_wandb
@@ -53,7 +54,8 @@ class ProgressTracker:
53
  def start_train(self, num_epochs: int):
54
  self.task_epoch = self.progress.add_task("[cyan]Epochs [white]| Loss | Box | DFL | BCE |", total=num_epochs)
55
 
56
- def start_one_epoch(self, num_batches, optimizer, epoch_idx):
 
57
  if self.use_wandb:
58
  lr_values = [params["lr"] for params in optimizer.param_groups]
59
  lr_names = ["bias", "norm", "conv"]
@@ -71,10 +73,10 @@ class ProgressTracker:
71
  loss_str += f" {loss_val:2.2f} |"
72
 
73
  self.progress.update(self.batch_task, advance=1, description=f"[green]Batches [white]{loss_str}")
 
74
 
75
  def finish_one_epoch(self):
76
  self.progress.remove_task(self.batch_task)
77
- self.progress.update(self.task_epoch, advance=1)
78
 
79
  def finish_train(self):
80
  self.wandb.finish()
 
22
  from rich.progress import BarColumn, Progress, TextColumn, TimeRemainingColumn
23
  from rich.table import Table
24
  from torch import Tensor
25
+ from torch.optim import Optimizer
26
 
27
  from yolo.config.config import Config, GeneralConfig, YOLOLayer
28
 
 
41
  self.progress = Progress(
42
  TextColumn("[progress.description]{task.description}"),
43
  BarColumn(bar_width=None),
44
+ TextColumn("{task.completed:.0f}/{task.total:.0f}"),
45
  TimeRemainingColumn(),
46
  )
47
  self.use_wandb = use_wandb
 
54
  def start_train(self, num_epochs: int):
55
  self.task_epoch = self.progress.add_task("[cyan]Epochs [white]| Loss | Box | DFL | BCE |", total=num_epochs)
56
 
57
+ def start_one_epoch(self, num_batches: int, optimizer: Optimizer, epoch_idx: int):
58
+ self.num_batches = num_batches
59
  if self.use_wandb:
60
  lr_values = [params["lr"] for params in optimizer.param_groups]
61
  lr_names = ["bias", "norm", "conv"]
 
73
  loss_str += f" {loss_val:2.2f} |"
74
 
75
  self.progress.update(self.batch_task, advance=1, description=f"[green]Batches [white]{loss_str}")
76
+ self.progress.update(self.task_epoch, advance=1 / self.num_batches)
77
 
78
  def finish_one_epoch(self):
79
  self.progress.remove_task(self.batch_task)
 
80
 
81
  def finish_train(self):
82
  self.wandb.finish()