From f35b3ad28f431fda5e5773699cd1c3c9450b8ee3 Mon Sep 17 00:00:00 2001 From: James Betker Date: Wed, 26 Aug 2020 08:44:22 -0600 Subject: [PATCH] Fix val behavior for ExtensibleTrainer --- codes/models/ExtensibleTrainer.py | 4 ++-- codes/models/archs/SwitchedResidualGenerator_arch.py | 5 ++++- codes/train.py | 2 +- codes/train2.py | 2 +- 4 files changed, 8 insertions(+), 5 deletions(-) diff --git a/codes/models/ExtensibleTrainer.py b/codes/models/ExtensibleTrainer.py index 501d6f67..919a945a 100644 --- a/codes/models/ExtensibleTrainer.py +++ b/codes/models/ExtensibleTrainer.py @@ -137,8 +137,8 @@ class ExtensibleTrainer(BaseModel): # Some models need to make parametric adjustments per-step. Do that here. for net in self.networks.values(): - if hasattr(net, "update_for_step"): - net.update_for_step(step, os.path.join(self.opt['path']['models'], "..")) + if hasattr(net.module, "update_for_step"): + net.module.update_for_step(step, os.path.join(self.opt['path']['models'], "..")) # Iterate through the steps, performing them one at a time. state = self.dstate diff --git a/codes/models/archs/SwitchedResidualGenerator_arch.py b/codes/models/archs/SwitchedResidualGenerator_arch.py index 58b02c86..377e0f61 100644 --- a/codes/models/archs/SwitchedResidualGenerator_arch.py +++ b/codes/models/archs/SwitchedResidualGenerator_arch.py @@ -101,7 +101,10 @@ def gather_2d(input, index): ind_nd = ind_nd.repeat((1, c)) ind_nd = ind_nd.unsqueeze(2) result = torch.gather(nodim, dim=2, index=ind_nd) - return result.squeeze() + result = result.squeeze() + if b == 1: + result = result.unsqueeze(0) + return result # Computes a linear latent by performing processing on the reference image and returning the filters of a single point, diff --git a/codes/train.py b/codes/train.py index 980d00c2..4c3e8761 100644 --- a/codes/train.py +++ b/codes/train.py @@ -215,7 +215,7 @@ def main(): logger.info(message) #### validation if opt['datasets'].get('val', None) and current_step % opt['train']['val_freq'] == 0: - if opt['model'] in ['sr', 'srgan', 'corruptgan', 'spsrgan'] and rank <= 0: # image restoration validation + if opt['model'] in ['sr', 'srgan', 'corruptgan', 'spsrgan', 'extensibletrainer'] and rank <= 0: # image restoration validation model.force_restore_swapout() val_batch_sz = 1 if 'batch_size' not in opt['datasets']['val'].keys() else opt['datasets']['val']['batch_size'] # does not support multi-GPU validation diff --git a/codes/train2.py b/codes/train2.py index 8dccf326..112cc83b 100644 --- a/codes/train2.py +++ b/codes/train2.py @@ -32,7 +32,7 @@ def init_dist(backend='nccl', **kwargs): def main(): #### options parser = argparse.ArgumentParser() - parser.add_argument('-opt', type=str, help='Path to option YAML file.', default='../options/train_mi1_nt_spsr_switched.yml') + parser.add_argument('-opt', type=str, help='Path to option YAML file.', default='../options/train_imgset_spsr_switched2_fullimgref.yml') parser.add_argument('--launcher', choices=['none', 'pytorch'], default='none', help='job launcher') parser.add_argument('--local_rank', type=int, default=0)