wandb integrate sample.py

Author: ayulockin

config/sr_sr3_16_128.json

@@ -17,7 +17,7 @@
  "name": "FFHQ",
  "mode": "HR", // whether need LR img
  "dataroot": "dataset/ffhq_16_128",
- "datatype": "img", //lmdb or img, path of img files
+ "datatype": "lmdb", //lmdb or img, path of img files
  "l_resolution": 16, // low resolution need to super_resolution
  "r_resolution": 128, // high resolution
  "batch_size": 4,
@@ -28,8 +28,8 @@
  "val": {
  "name": "CelebaHQ",
  "mode": "LRHR",
- "dataroot": "dataset/ffhq_16_128",
- "datatype": "img", //lmdb or img, path of img files
+ "dataroot": "dataset/celebahq_16_128",
+ "datatype": "lmdb", //lmdb or img, path of img files
  "l_resolution": 16,
  "r_resolution": 128,
  "data_len": 50 // data length in validation 

sample.py

@@ -5,9 +5,11 @@
 import logging
 import core.logger as Logger
 import core.metrics as Metrics
+from core.wandb_logger import WandbLogger
 from tensorboardX import SummaryWriter
 import os
 import numpy as np
+import wandb
 
 if __name__ == "__main__":
  parser = argparse.ArgumentParser()
@@ -17,6 +19,9 @@
  help='Run either train(training) or val(generation)', default='train')
  parser.add_argument('-gpu', '--gpu_ids', type=str, default=None)
  parser.add_argument('-debug', '-d', action='store_true')
+ parser.add_argument('-enable_wandb', action='store_true')
+ parser.add_argument('-log_wandb_ckpt', action='store_true')
+ parser.add_argument('-log_eval', action='store_true')
 
  # parse configs
  args = parser.parse_args()
@@ -35,6 +40,16 @@
  logger.info(Logger.dict2str(opt))
  tb_logger = SummaryWriter(log_dir=opt['path']['tb_logger'])
 
+ # Initialize WandbLogger
+ if opt['enable_wandb']:
+ wandb_logger = WandbLogger(opt)
+ # wandb.define_metric('validation/val_step')
+ # wandb.define_metric('epoch')
+ # wandb.define_metric("validation/*", step_metric="val_step")
+ val_step = 0
+ else:
+ wandb_logger = None
+
  # dataset
  for phase, dataset_opt in opt['datasets'].items():
  if phase == 'train' and args.phase != 'val':
@@ -78,6 +93,9 @@
  tb_logger.add_scalar(k, v, current_step)
  logger.info(message)
 
+ if wandb_logger:
+ wandb_logger.log_metrics(logs)
+
  # validation
  if current_step % opt['train']['val_freq'] == 0:
  result_path = '{}/{}'.format(opt['path']
@@ -100,12 +118,20 @@
  'Iter_{}'.format(current_step),
  np.transpose(sample_img, [2, 0, 1]),
  idx)
+
+ if wandb_logger:
+ wandb_logger.log_image(f'validation_{idx}', sample_img)
+
  diffusion.set_new_noise_schedule(
  opt['model']['beta_schedule']['train'], schedule_phase='train')
 
  if current_step % opt['train']['save_checkpoint_freq'] == 0:
  logger.info('Saving models and training states.')
  diffusion.save_network(current_epoch, current_step)
+
+ if wandb_logger and opt['log_wandb_ckpt']:
+ wandb_logger.log_checkpoint(current_epoch, current_step)
+
  # save model
  logger.info('End of training.')
  else: