1import time 

2import wandb 

3from fastai.callback.wandb import WandbCallback 

4 

5 

6class TorchAppWandbCallback(WandbCallback): 

7 def __init__(self, app, project_name=None, **kwargs): 

8 self.app = app 

9 if project_name is None: 

10 project_name = app.project_name() 

11 

12 self.run = wandb.init(project=project_name, reinit=True, **kwargs) 

13 

14 super().__init__() 

15 

16 if hasattr(app, 'training_kwargs'): 

17 wandb.config.update(app.training_kwargs) 

18 

19 def after_epoch(self): 

20 super().after_epoch() 

21 

22 # Record the length of time of each epoch 

23 wandb.log( 

24 {"time": time.time() - self.recorder.start_epoch}, 

25 step=self._wandb_step, 

26 )