def init_weights(self, pretrained=None):
"""Init weights for the model.
Args:
pretrained (str, optional): Path for pretrained weights. If given
None, pretrained weights will not be loaded. Defaults to None.
"""
if isinstance(pretrained, str):
logger = get_root_logger()
load_checkpoint(self, pretrained, strict=False, logger=logger)
elif pretrained is None:
for m in self.modules():
if isinstance(m, nn.Conv2d):
kaiming_init(m)
elif isinstance(m, (_BatchNorm, nn.GroupNorm)):
constant_init(m, 1)
if self.zero_init_residual:
for m in self.modules():
if isinstance(m, Bottleneck):
constant_init(m.norm3, 0)
elif isinstance(m, BasicBlock):
constant_init(m.norm2, 0)
else:
raise TypeError('pretrained must be a str or None')
def init_weights(self, pretrained=None, strict=True):
"""Init weights for models.
Args:
pretrained (str, optional): Path for pretrained weights. If given
None, pretrained weights will not be loaded. Defaults to None.
strict (boo, optional): Whether strictly load the pretrained model.
Defaults to True.
"""
if isinstance(pretrained, str):
logger = get_root_logger()
load_checkpoint(self, pretrained, strict=strict, logger=logger)
elif pretrained is None:
if self.with_tsa:
for module in [
self.fusion.feat_fusion, self.fusion.spatial_attn1,
self.fusion.spatial_attn2, self.fusion.spatial_attn3,
self.fusion.spatial_attn4, self.fusion.spatial_attn_l1,
self.fusion.spatial_attn_l2,
self.fusion.spatial_attn_l3,
self.fusion.spatial_attn_add1
]:
kaiming_init(module.conv,
a=0.1,
mode='fan_out',
nonlinearity='leaky_relu',
bias=0,
distribution='uniform')
else:
raise TypeError(f'"pretrained" must be a str or None. '
f'But received {type(pretrained)}.')
def main():
parser = argparse.ArgumentParser(description='Benchmark dataloading')
parser.add_argument('config', help='train config file path')
args = parser.parse_args()
cfg = Config.fromfile(args.config)
# init logger before other steps
logger = get_root_logger()
logger.info(f'Config: {cfg.text}')
dataset = build_dataset(cfg.data.train)
data_loaders = [
build_dataloader(ds,
cfg.data.samples_per_gpu,
cfg.data.workers_per_gpu,
dist=False,
drop_last=cfg.data.get('drop_last', False),
seed=0) for ds in dataset
]
# Start progress bar after first 5 batches
prog_bar = mmcv.ProgressBar(len(dataset) - 5 * cfg.data.samples_per_gpu,
start=False)
for data_loader in data_loaders:
for i, data in enumerate(data_loader):
if i == 5:
prog_bar.start()
for _ in data['imgs']:
if i < 5:
continue
prog_bar.update()
def test_visual_hook():
with pytest.raises(
AssertionError), tempfile.TemporaryDirectory() as tmpdir:
VisualizationHook(tmpdir, [1, 2, 3])
test_dataset = ExampleDataset()
test_dataset.evaluate = MagicMock(return_value=dict(test='success'))
img = torch.zeros((1, 3, 10, 10))
img[:, :, 2:9, :] = 1.
model = ExampleModel()
data_loader = DataLoader(test_dataset,
batch_size=1,
sampler=None,
num_workers=0,
shuffle=False)
with tempfile.TemporaryDirectory() as tmpdir:
visual_hook = VisualizationHook(tmpdir, ['img'],
interval=8,
rerange=False)
runner = mmcv.runner.IterBasedRunner(model=model,
optimizer=None,
work_dir=tmpdir,
logger=get_root_logger())
runner.register_hook(visual_hook)
runner.run([data_loader], [('train', 10)], 10)
img_saved = mmcv.imread(osp.join(tmpdir, 'iter_8.png'),
flag='unchanged')
np.testing.assert_almost_equal(img_saved,
img[0].permute(1, 2, 0) * 255)
with tempfile.TemporaryDirectory() as tmpdir:
visual_hook = VisualizationHook(tmpdir, ['img'],
interval=8,
rerange=True)
runner = mmcv.runner.IterBasedRunner(model=model,
optimizer=None,
work_dir=tmpdir,
logger=get_root_logger())
runner.register_hook(visual_hook)
runner.run([data_loader], [('train', 10)], 10)
img_saved = mmcv.imread(osp.join(tmpdir, 'iter_8.png'),
flag='unchanged')
assert osp.exists(osp.join(tmpdir, 'iter_8.png'))
def init_weights(self, model, pretrained):
"""Init weights.
Args:
model (nn.Module): Models to be inited.
pretrained (str): Path for pretrained weights.
"""
logger = get_root_logger()
load_checkpoint(model, pretrained, logger=logger)
def init_weights(self, pretrained=None):
if isinstance(pretrained, str):
logger = get_root_logger()
load_checkpoint(self, pretrained, strict=False, logger=logger)
elif pretrained is None:
for m in self.modules():
if isinstance(m, nn.Conv2d):
xavier_init(m)
elif isinstance(m, nn.BatchNorm2d):
constant_init(m, 1)
def init_weights(self, pretrained=None, strict=True):
"""Init weights for models.
Args:
pretrained (str, optional): Path for pretrained weights. If given
None, pretrained weights will not be loaded. Defaults to None.
strict (boo, optional): Whether strictly load the pretrained model.
Defaults to True.
"""
if isinstance(pretrained, str):
logger = get_root_logger()
load_checkpoint(self, pretrained, strict=strict, logger=logger)
elif pretrained is not None:
raise TypeError(f'"pretrained" must be a str or None. '
f'But received {type(pretrained)}.')
def init_weights(self, pretrained=None):
"""Init weights for models.
Args:
pretrained (str, optional): Path for pretrained weights. If given
None, pretrained weights will not be loaded. Defaults to None.
"""
if isinstance(pretrained, str):
logger = get_root_logger()
load_checkpoint(self, pretrained, strict=False, logger=logger)
elif pretrained is None:
# Here, we just use the default initialization in `ConvModule`.
pass
else:
raise TypeError('pretrained must be a str or None')
def init_weights(self, pretrained=None):
"""Initialize weights for the model.
Args:
pretrained (str, optional): Path for pretrained weights. If given
None, pretrained weights will not be loaded. Default: None.
"""
if isinstance(pretrained, str):
logger = get_root_logger()
load_checkpoint(self, pretrained, strict=False, logger=logger)
elif pretrained is None:
generation_init_weights(
self, init_type=self.init_type, init_gain=self.init_gain)
else:
raise TypeError("'pretrained' must be a str or None. "
f'But received {type(pretrained)}.')
def init_weights(self, pretrained=None):
"""Init weights for the model.
Args:
pretrained (str, optional): Path for pretrained weights. If given
None, pretrained weights will not be loaded. Defaults to None.
"""
if isinstance(pretrained, str):
logger = get_root_logger()
load_checkpoint(self, pretrained, strict=False, logger=logger)
elif pretrained is None:
for m in self.modules():
if isinstance(m, nn.Conv2d):
xavier_init(m)
elif isinstance(m, nn.BatchNorm2d):
constant_init(m, 1)
def __init__(self, pretrained):
super().__init__()
self.basic_module = nn.ModuleList(
[SPyNetBasicModule() for _ in range(6)])
if isinstance(pretrained, str):
logger = get_root_logger()
load_checkpoint(self, pretrained, strict=True, logger=logger)
elif pretrained is not None:
raise TypeError('[pretrained] should be str or None, '
f'but got {type(pretrained)}.')
self.register_buffer(
'mean',
torch.Tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))
self.register_buffer(
'std',
torch.Tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))
def init_weights(self, pretrained=None):
"""Init weights for models.
Args:
pretrained (str, optional): Path for pretrained weights. If given
None, pretrained weights will not be loaded. Defaults to None.
"""
if isinstance(pretrained, str):
logger = get_root_logger()
load_checkpoint(self, pretrained, strict=False, logger=logger)
elif pretrained is None:
for m in self.modules():
# Here, we only initialize the module with fc layer since the
# conv and norm layers has been intialized in `ConvModule`.
if isinstance(m, nn.Linear):
nn.init.normal_(m.weight.data, 0.0, 0.02)
nn.init.constant_(m.bias.data, 0.0)
else:
raise TypeError('pretrained must be a str or None')
def init_weights(self, pretrained=None):
"""Init weights for models.
Args:
pretrained (str, optional): Path for pretrained weights. If given
None, pretrained weights will not be loaded. Defaults to None.
"""
if isinstance(pretrained, str):
logger = get_root_logger()
load_checkpoint(self, pretrained, strict=False, logger=logger)
elif pretrained is None:
for m in self.modules():
if isinstance(m, nn.Linear):
normal_init(m, 0, std=0.02)
elif isinstance(m, nn.Conv2d):
normal_init(m, 0.0, std=0.02)
else:
raise TypeError('pretrained must be a str or None but got'
f'{type(pretrained)} instead.')
def train_model(model,
dataset,
cfg,
distributed=False,
validate=False,
timestamp=None,
meta=None):
"""Train model entry function.
Args:
model (nn.Module): The model to be trained.
dataset (:obj:`Dataset`): Train dataset.
cfg (dict): The config dict for training.
distributed (bool): Whether to use distributed training.
Default: False.
validate (bool): Whether to do evaluation. Default: False.
timestamp (str | None): Local time for runner. Default: None.
meta (dict | None): Meta dict to record some important information.
Default: None
"""
logger = get_root_logger(log_level=cfg.log_level)
# start training
if distributed:
_dist_train(
model,
dataset,
cfg,
validate=validate,
logger=logger,
timestamp=timestamp,
meta=meta)
else:
_non_dist_train(
model,
dataset,
cfg,
validate=validate,
logger=logger,
timestamp=timestamp,
meta=meta)
def init_weights(self, pretrained=None, strict=True):
"""Init weights for models.
Args:
pretrained (str, optional): Path for pretrained weights. If given
None, pretrained weights will not be loaded. Defaults to None.
strict (boo, optional): Whether strictly load the pretrained model.
Defaults to True.
"""
if isinstance(pretrained, str):
logger = get_root_logger()
load_checkpoint(self, pretrained, strict=strict, logger=logger)
elif pretrained is None:
# Initialization methods like `kaiming_init` are for VGG-style
# modules. For modules with residual paths, using smaller std is
# better for stability and performance. There is a global residual
# path in MSRResNet and empirically we use 0.1. See more details in
# "ESRGAN: Enhanced Super-Resolution Generative Adversarial
# Networks"
for m in [self.conv_first, self.conv_hr, self.conv_last]:
default_init_weights(m, 0.1)
else:
raise TypeError(f'"pretrained" must be a str or None. '
f'But received {type(pretrained)}.')
def init_weights(self, pretrained=None, strict=True):
"""Init weights for models.
Args:
pretrained (str, optional): Path for pretrained weights. If given
None, pretrained weights will not be loaded. Defaults to None.
strict (boo, optional): Whether strictly load the pretrained model.
Defaults to True.
"""
if isinstance(pretrained, str):
logger = get_root_logger()
load_checkpoint(self, pretrained, strict=strict, logger=logger)
elif pretrained is None:
# Use smaller std for better stability and performance. We
# use 0.1. See more details in "ESRGAN: Enhanced Super-Resolution
# Generative Adversarial Networks"
for m in [
self.conv_first, self.conv_body, self.conv_up1,
self.conv_up2, self.conv_hr, self.conv_last
]:
default_init_weights(m, 0.1)
else:
raise TypeError(f'"pretrained" must be a str or None. '
f'But received {type(pretrained)}.')
def __init__(self,
in_channels=3,
out_channel=3,
mid_channels=64,
num_frames=5,
deform_groups=8,
num_blocks_extraction=5,
num_blocks_reconstruction=10,
center_frame_idx=2,
with_tsa=True,
pretrained=None):
super().__init__()
self.center_frame_idx = center_frame_idx
self.with_tsa = with_tsa
act_cfg = dict(type='LeakyReLU', negative_slope=0.1)
self.conv_first = nn.Conv2d(in_channels, mid_channels, 3, 1, 1)
self.feature_extraction = make_layer(ResidualBlockNoBN,
num_blocks_extraction,
mid_channels=mid_channels)
# generate pyramid features
self.feat_l2_conv1 = ConvModule(mid_channels,
mid_channels,
3,
2,
1,
act_cfg=act_cfg)
self.feat_l2_conv2 = ConvModule(mid_channels,
mid_channels,
3,
1,
1,
act_cfg=act_cfg)
self.feat_l3_conv1 = ConvModule(mid_channels,
mid_channels,
3,
2,
1,
act_cfg=act_cfg)
self.feat_l3_conv2 = ConvModule(mid_channels,
mid_channels,
3,
1,
1,
act_cfg=act_cfg)
# pcd alignment
self.pcd_alignment = PCDAlignment(mid_channels=mid_channels,
deform_groups=deform_groups)
# fusion
if self.with_tsa:
self.fusion = TSAFusion(mid_channels=mid_channels,
num_frames=num_frames,
center_frame_idx=self.center_frame_idx)
else:
self.fusion = nn.Conv2d(num_frames * mid_channels, mid_channels, 1,
1)
# activation function
self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
if isinstance(pretrained, str):
logger = get_root_logger()
load_checkpoint(self, pretrained, strict=True, logger=logger)
elif pretrained is not None:
raise TypeError(f'"pretrained" must be a str or None. '
f'But received {type(pretrained)}.')
def main():
args = parse_args()
cfg = Config.fromfile(args.config)
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
# update configs according to CLI args
if args.work_dir is not None:
cfg.work_dir = args.work_dir
if args.resume_from is not None:
cfg.resume_from = args.resume_from
cfg.gpus = args.gpus
if args.autoscale_lr:
# apply the linear scaling rule (https://arxiv.org/abs/1706.02677)
cfg.optimizer['lr'] = cfg.optimizer['lr'] * cfg.gpus / 8
# init distributed env first, since logger depends on the dist info.
if args.launcher == 'none':
distributed = False
else:
distributed = True
init_dist(args.launcher, **cfg.dist_params)
# create work_dir
mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir))
# init the logger before other steps
timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
log_file = osp.join(cfg.work_dir, f'{timestamp}.log')
logger = get_root_logger(log_file=log_file, log_level=cfg.log_level)
# log env info
env_info_dict = collect_env.collect_env()
env_info = '\n'.join([f'{k}: {v}' for k, v in env_info_dict.items()])
dash_line = '-' * 60 + '\n'
logger.info('Environment info:\n' + dash_line + env_info + '\n' +
dash_line)
# log some basic info
logger.info('Distributed training: {}'.format(distributed))
logger.info('mmedit Version: {}'.format(__version__))
logger.info('Config:\n{}'.format(cfg.text))
# set random seeds
if args.seed is not None:
logger.info('Set random seed to {}, deterministic: {}'.format(
args.seed, args.deterministic))
set_random_seed(args.seed, deterministic=args.deterministic)
cfg.seed = args.seed
model = build_model(cfg.model,
train_cfg=cfg.train_cfg,
test_cfg=cfg.test_cfg)
datasets = [build_dataset(cfg.data.train)]
if len(cfg.workflow) == 2:
val_dataset = copy.deepcopy(cfg.data.val)
val_dataset.pipeline = cfg.data.train.pipeline
datasets.append(build_dataset(val_dataset))
if cfg.checkpoint_config is not None:
# save version, config file content and class names in
# checkpoints as meta data
cfg.checkpoint_config.meta = dict(
mmedit_version=__version__,
config=cfg.text,
)
# meta information
meta = dict()
if cfg.get('exp_name', None) is None:
cfg['exp_name'] = osp.splitext(osp.basename(cfg.work_dir))[0]
meta['exp_name'] = cfg.exp_name
meta['mmedit Version'] = __version__
meta['seed'] = args.seed
meta['env_info'] = env_info
# add an attribute for visualization convenience
train_model(model,
datasets,
cfg,
distributed=distributed,
validate=(not args.no_validate),
timestamp=timestamp,
meta=meta)