Spaces:
Runtime error
Runtime error
File size: 11,188 Bytes
5e88f62 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 |
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
"""
MaskFormer Training Script.
This script is a simplified version of the training script in detectron2/tools.
"""
import copy
import itertools
import json
import logging
import os
import sys
from typing import Any, Dict, List, Set
import detectron2.utils.comm as comm
import torch
import wandb
from detectron2.config import get_cfg, CfgNode
from detectron2.engine import DefaultTrainer, default_setup
from detectron2.projects.deeplab import add_deeplab_config, build_lr_scheduler
from detectron2.solver.build import maybe_add_gradient_clipping
from detectron2.utils.file_io import PathManager
from detectron2.utils.logger import setup_logger
import utils
# MaskFormer
from config import add_gwm_config
logger = logging.getLogger('gwm')
class Trainer(DefaultTrainer):
"""
Extension of the Trainer class adapted to DETR.
"""
@classmethod
def build_evaluator(cls, cfg, dataset_name):
pass
@classmethod
def build_lr_scheduler(cls, cfg, optimizer):
"""
It now calls :func:`detectron2.solver.build_lr_scheduler`.
Overwrite it if you'd like a different scheduler.
"""
return build_lr_scheduler(cfg, optimizer)
@classmethod
def build_optimizer(cls, cfg, model):
weight_decay_norm = cfg.SOLVER.WEIGHT_DECAY_NORM
weight_decay_embed = cfg.SOLVER.WEIGHT_DECAY_EMBED
defaults = {}
defaults["lr"] = cfg.SOLVER.BASE_LR
defaults["weight_decay"] = cfg.SOLVER.WEIGHT_DECAY
norm_module_types = (
torch.nn.BatchNorm1d,
torch.nn.BatchNorm2d,
torch.nn.BatchNorm3d,
torch.nn.SyncBatchNorm,
# NaiveSyncBatchNorm inherits from BatchNorm2d
torch.nn.GroupNorm,
torch.nn.InstanceNorm1d,
torch.nn.InstanceNorm2d,
torch.nn.InstanceNorm3d,
torch.nn.LayerNorm,
torch.nn.LocalResponseNorm,
)
params: List[Dict[str, Any]] = []
memo: Set[torch.nn.parameter.Parameter] = set()
for module_name, module in model.named_modules():
for module_param_name, value in module.named_parameters(recurse=False):
if not value.requires_grad:
continue
# Avoid duplicating parameters
if value in memo:
continue
memo.add(value)
hyperparams = copy.copy(defaults)
if "backbone" in module_name:
hyperparams["lr"] = hyperparams["lr"] * cfg.SOLVER.BACKBONE_MULTIPLIER
if (
"relative_position_bias_table" in module_param_name
or "absolute_pos_embed" in module_param_name
):
hyperparams["weight_decay"] = 0.0
if isinstance(module, norm_module_types):
hyperparams["weight_decay"] = weight_decay_norm
if isinstance(module, torch.nn.Embedding):
hyperparams["weight_decay"] = weight_decay_embed
params.append({"params": [value], **hyperparams})
def maybe_add_full_model_gradient_clipping(optim):
# detectron2 doesn't have full model gradient clipping now
clip_norm_val = cfg.SOLVER.CLIP_GRADIENTS.CLIP_VALUE
enable = (
cfg.SOLVER.CLIP_GRADIENTS.ENABLED
and cfg.SOLVER.CLIP_GRADIENTS.CLIP_TYPE == "full_model"
and clip_norm_val > 0.0
)
class FullModelGradientClippingOptimizer(optim):
def step(self, closure=None):
all_params = itertools.chain(*[x["params"] for x in self.param_groups])
torch.nn.utils.clip_grad_norm_(all_params, clip_norm_val)
super().step(closure=closure)
return FullModelGradientClippingOptimizer if enable else optim
optimizer_type = cfg.SOLVER.OPTIMIZER
if optimizer_type == "SGD":
optimizer = maybe_add_full_model_gradient_clipping(torch.optim.SGD)(
params, cfg.SOLVER.BASE_LR, momentum=cfg.SOLVER.MOMENTUM
)
elif optimizer_type == "ADAMW":
optimizer = maybe_add_full_model_gradient_clipping(torch.optim.AdamW)(
params, cfg.SOLVER.BASE_LR
)
elif optimizer_type == "RMSProp":
optimizer = maybe_add_full_model_gradient_clipping(torch.optim.RMSprop)(
params, cfg.SOLVER.BASE_LR
)
else:
raise NotImplementedError(f"no optimizer type {optimizer_type}")
if not cfg.SOLVER.CLIP_GRADIENTS.CLIP_TYPE == "full_model":
optimizer = maybe_add_gradient_clipping(cfg, optimizer)
return optimizer
def setup(args):
"""
Create configs and perform basic setups.
"""
wandb_inited = False
if 'CONFIG_FILE' in args.opts and not args.wandb_sweep_mode:
logger.warning(
f"Found CONFIG_FILE key in OPT args and using {args.opts[args.opts.index('CONFIG_FILE') + 1]} instead of {args.config_file}")
args.config_file = args.opts[args.opts.index('CONFIG_FILE') + 1]
else:
cfg = get_cfg()
add_gwm_config(cfg)
wandb_basedir = cfg.WANDB.BASEDIR
cfg_dict = CfgNode.load_yaml_with_base(args.config_file, allow_unsafe=True)
if 'WANDB' in cfg_dict and 'BASEDIR' in cfg_dict['WANDB']:
wandb_basedir = cfg_dict['WANDB']['BASEDIR']
if 'CONFIG_FILE' in cfg_dict and cfg_dict['CONFIG_FILE'] is not None:
logger.warning(
f"Found CONFIG_FILE key in the config.yaml file and using {cfg_dict['CONFIG_FILE']} instead of {args.config_file}")
args.config_file = cfg_dict['CONFIG_FILE']
if args.wandb_sweep_mode:
if PathManager.isfile('wandb.yaml'):
wandb_cfg = CfgNode.load_yaml_with_base('wandb.yaml', allow_unsafe=False)
wandb.init(project=wandb_cfg['PROJECT'], entity=wandb_cfg['USER'], dir=wandb_basedir)
wandb_inited = True
if wandb.run.sweep_id: # sweep active
sweep_dict = dict(wandb.config)
if 'CONFIG_FILE' in sweep_dict:
args.config_file = sweep_dict['CONFIG_FILE']
logger.warning(f"Loading CONFIG_FILE as set in sweep config: {args.config_file}")
elif 'CONFIG_FILE' in args.opts:
args.config_file = args.opts[args.opts.index('CONFIG_FILE') + 1]
logger.warning(f"Loading CONFIG_FILE as set in the optional arguments: {args.config_file}")
if 'GWM.MODEL' in args.opts and not args.wandb_sweep_mode:
logger.warning(
"It is advised to not set GWM.MODEL in OPT args and instead set it in the config.yaml file")
model = args.opts[args.opts.index('GWM.MODEL') + 1]
else:
cfg = get_cfg()
add_gwm_config(cfg)
model = cfg.GWM.MODEL
wandb_basedir = cfg.WANDB.BASEDIR
cfg_dict = CfgNode.load_yaml_with_base(args.config_file, allow_unsafe=True)
if 'GWM' in cfg_dict and 'MODEL' in cfg_dict['GWM']:
model = cfg_dict['GWM']['MODEL']
if 'WANDB' in cfg_dict and 'BASEDIR' in cfg_dict['WANDB']:
wandb_basedir = cfg_dict['WANDB']['BASEDIR']
if args.wandb_sweep_mode:
if PathManager.isfile('wandb.yaml'):
if not wandb_inited:
wandb_cfg = CfgNode.load_yaml_with_base('wandb.yaml', allow_unsafe=False)
wandb.init(project=wandb_cfg['PROJECT'], entity=wandb_cfg['USER'], dir=wandb_basedir)
wandb_inited = True
if args.wandb_sweep_mode:
sweep_dict = dict(wandb.config)
if 'GWM.MODEL' in sweep_dict:
logger.warning(
"It is advised to not set GWM.MODEL in sweep config and instead set it in the config.yaml file")
model = sweep_dict['GWM.MODEL']
elif 'GWM.MODEL' in args.opts:
logger.warning(
"It is advised to not set GWM.MODEL in optional arguments and instead set it in the config.yaml file")
model = args.opts[args.opts.index('GWM.MODEL') + 1]
cfg = get_cfg()
# for poly lr schedule
add_deeplab_config(cfg)
if model == "MASKFORMER":
from mask_former import add_mask_former_config
add_mask_former_config(cfg)
else:
logger.error(f'Unknown Model: {model}. Exiting..')
sys.exit(0)
add_gwm_config(cfg)
cfg.merge_from_file(args.config_file)
cfg.merge_from_list(args.opts)
cfg.WANDB.ENABLE = (cfg.WANDB.ENABLE or args.wandb_sweep_mode) and not args.eval_only
datestring = utils.log.get_datestring_for_the_run()
if cfg.WANDB.ENABLE:
if PathManager.isfile('wandb.yaml'):
if not wandb_inited:
wandb_cfg = CfgNode.load_yaml_with_base('wandb.yaml', allow_unsafe=False)
wandb.init(project=wandb_cfg['PROJECT'], entity=wandb_cfg['USER'], dir=cfg.WANDB.BASEDIR)
if args.wandb_sweep_mode: # sweep active
sweep_list = [(k, v) for k, v in dict(wandb.config).items()]
sweep_list = [item for items in sweep_list for item in items]
cfg.merge_from_list(sweep_list)
if cfg.LOG_ID is not None:
api = wandb.Api()
run = api.run(path=f"{wandb_cfg['USER']}/{wandb_cfg['PROJECT']}/{wandb.run.id}")
run.name = f'{cfg.LOG_ID}/{datestring}-{wandb.run.id}'
run.save()
else:
logger.error("W&B config file 'src/wandb.yaml' does not exist!")
cfg.WANDB.ENABLE = False
if args.resume_path:
cfg.OUTPUT_DIR = "/".join(args.resume_path.split('/')[:-2]) # LOG_ID/datestring/checkpoints/checkpoints.pth
if args.eval_only:
cfg.OUTPUT_DIR = os.path.join(cfg.OUTPUT_DIR, 'eval', datestring)
else:
if cfg.LOG_ID and not cfg.SLURM:
cfg.OUTPUT_DIR = os.path.join(cfg.OUTPUT_BASEDIR, cfg.LOG_ID)
else:
cfg.OUTPUT_DIR = cfg.OUTPUT_BASEDIR
if args.eval_only:
cfg.OUTPUT_DIR = None
else:
cfg.OUTPUT_DIR = os.path.join(cfg.OUTPUT_DIR, datestring)
os.makedirs(f'{cfg.OUTPUT_DIR}/checkpoints', exist_ok=True)
if cfg.WANDB.ENABLE:
wandb.config.update(cfg, allow_val_change=True)
if cfg.GWM.LOSS == 'OG':
cfg.FLAGS.EXTENDED_FLOW_RECON_VIS = False
cfg.FLAGS.COMP_NLL_FOR_GT = False
cfg.freeze()
default_setup(cfg, args)
# Setup logger for "gwm" module
setup_logger(output=f'{cfg.OUTPUT_DIR}/main.log', distributed_rank=comm.get_rank(), name="gwm")
with open(f'{cfg.OUTPUT_DIR}/args.json', 'w') as f:
json.dump(args.__dict__, f, indent=2)
return cfg
|