Merge pull request #246 from rom1504/sample_per_sec

Sample per sec

Author: lucidrains

train_dalle.py

@@ -1,5 +1,6 @@
 import argparse
 from pathlib import Path
+import time
 
 import torch
 import wandb  # Quit early if user doesn't have wandb installed.
@@ -381,6 +382,8 @@ def save_model(path):
     if data_sampler:
         data_sampler.set_epoch(epoch)
     for i, (text, images) in enumerate(distr_dl):
+        if i % 10 == 0 and distr_backend.is_root_worker():
+            t = time.time()
         if args.fp16:
             images = images.half()
         text, images = map(lambda t: t.cuda(), (text, images))
@@ -433,6 +436,11 @@ def save_model(path):
                     log['image'] = wandb.Image(image, caption=decoded_text)
 
 
+        if i % 10 == 9 and distr_backend.is_root_worker():
+            sample_per_sec = BATCH_SIZE * 10 / (time.time() - t)
+            log["sample_per_sec"] = sample_per_sec
+            print(epoch, i, f'sample_per_sec - {sample_per_sec}')
+
         if distr_backend.is_root_worker():
             wandb.log(log)