-
Notifications
You must be signed in to change notification settings - Fork 385
/
opt.py
executable file
·177 lines (153 loc) · 8.67 KB
/
opt.py
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
import argparse
import yaml
def get_opts():
parser = argparse.ArgumentParser()
# General Setttings
parser.add_argument('--root_dir', type=str, default='Batman_masked_frames',
help='root directory of dataset')
parser.add_argument('--canonical_dir', type=str, default=None,
help='directory of canonical dataset')
# support multiple mask as input (each mask has different deformation fields)
parser.add_argument('--mask_dir', nargs="+", type=str, default=None,
help='mask of the dataset')
parser.add_argument('--flow_dir', type=str,
default=None,
help='masks of dataset')
parser.add_argument('--dataset_name', type=str, default='video',
choices=['video'],
help='which dataset to train/val')
parser.add_argument('--img_wh', nargs="+", type=int, default=[842, 512],
help='resolution (img_w, img_h) of the full image')
parser.add_argument('--canonical_wh', nargs="+", type=int, default=None,
help='default same as the img_wh, can be set to a larger range to include more content')
parser.add_argument('--ref_idx', type=int, default=None,
help='manually select a frame as reference (for rigid movement)')
# Deformation Setting
parser.add_argument('--encode_w', default=False, action="store_true",
help='whether to apply warping')
# Training Setttings
parser.add_argument('--batch_size', type=int, default=1,
help='batch size')
parser.add_argument('--num_steps', type=int, default=10000,
help='number of training epochs')
parser.add_argument('--valid_iters', type=int, default=30,
help='valid iters for each epoch')
parser.add_argument('--valid_batches', type=int, default=0,
help='valid batches for each valid process')
parser.add_argument('--save_model_iters', type=int, default=5000,
help='iterations to save the models')
parser.add_argument('--gpus', nargs="+", type=int, default=[0],
help='gpu devices')
# Test Setttings
parser.add_argument('--test', default=False, action="store_true",
help='whether to disable identity')
# Model Save and Load
parser.add_argument('--ckpt_path', type=str, default=None,
help='pretrained checkpoint to load (including optimizers, etc)')
parser.add_argument('--prefixes_to_ignore', nargs='+', type=str, default=['loss'],
help='the prefixes to ignore in the checkpoint state dict')
parser.add_argument('--weight_path', type=str, default=None,
help='pretrained model weight to load (do not load optimizers, etc)')
parser.add_argument('--model_save_path', type=str, default='ckpts',
help='save checkpoint to')
parser.add_argument('--log_save_path', type=str, default='logs',
help='save log to')
parser.add_argument('--exp_name', type=str, default='exp',
help='experiment name')
# Optimize Settings
parser.add_argument('--optimizer', type=str, default='adam',
help='optimizer type',
choices=['sgd', 'adam', 'radam', 'ranger'])
parser.add_argument('--lr', type=float, default=5e-4,
help='learning rate')
parser.add_argument('--momentum', type=float, default=0.9,
help='learning rate momentum')
parser.add_argument('--weight_decay', type=float, default=0,
help='weight decay')
parser.add_argument('--lr_scheduler', type=str, default='steplr',
help='scheduler type',
choices=['steplr', 'cosine', 'poly', 'exponential'])
#### params for steplr ####
parser.add_argument('--decay_step', nargs='+', type=int,
default=[2500, 5000, 7500],
help='scheduler decay step')
parser.add_argument('--decay_gamma', type=float, default=0.5,
help='learning rate decay amount')
#### params for warmup, only applied when optimizer == 'sgd' or 'adam'
parser.add_argument('--warmup_multiplier', type=float, default=1.0,
help='lr is multiplied by this factor after --warmup_epochs')
parser.add_argument('--warmup_epochs', type=int, default=0,
help='Gradually warm-up(increasing) learning rate in optimizer')
##### annealed positional encoding ######
parser.add_argument('--annealed', default=False, action="store_true",
help='whether to apply annealed positional encoding (Only in the warping field)')
parser.add_argument('--annealed_begin_step', type=int, default=0,
help='annealed step to begin for positional encoding')
parser.add_argument('--annealed_step', type=int, default=5000,
help='maximum annealed step for positional encoding')
##### Additional losses ######
parser.add_argument('--flow_loss', type=float, default=None,
help='optical flow loss weight')
parser.add_argument('--bg_loss', type=float, default=None,
help='regularize the rest part of each object ')
parser.add_argument('--grad_loss', type=float, default=0.1,
help='image gradient loss weight')
parser.add_argument('--flow_step', type=int, default=-1,
help='Step to begin to perform flow loss.')
parser.add_argument('--ref_step', type=int, default=-1,
help='Step to stop reference frame loss.')
parser.add_argument('--self_bg', type=bool_parser, default=False,
help='Whether to use self background as bg loss.')
##### Special cases: for black-dominated images
parser.add_argument('--sigmoid_offset', type=float, default=0,
help='whether to process balck-dominated images.')
# Other miscellaneous settings.
parser.add_argument('--save_deform', type=bool_parser, default=False,
help='Whether to save deformation field or not.')
parser.add_argument('--save_video', type=bool_parser, default=True,
help='Whether to save video or not.')
parser.add_argument('--fps', type=int, default=30,
help='FPS of the saved video.')
# Network settings for PE.
parser.add_argument('--deform_D', type=int, default=6,
help='The depth of deformation field MLP.')
parser.add_argument('--deform_W', type=int, default=128,
help='The width of deformation field MLP.')
parser.add_argument('--vid_D', type=int, default=8,
help='The depth of implicit video MLP.')
parser.add_argument('--vid_W', type=int, default=256,
help='The width of implicit video MLP.')
parser.add_argument('--N_vocab_w', type=int, default=200,
help='number of vocabulary for warp code in the dataset for nn.Embedding')
parser.add_argument('--N_w', type=int, default=8,
help='embeddings size for warping')
parser.add_argument('--N_xyz_w', nargs="+", type=int, default=[8, 8],
help='positional encoding frequency of deformation field')
# Network settings for Hash, please see details in configs/hash.json
parser.add_argument('--vid_hash', type=bool_parser, default=False,
help='Whether to use hash encoding in implicit video system.')
parser.add_argument('--deform_hash', type=bool_parser, default=False,
help='Whether to use hash encoding in deformation field.')
# Config files
parser.add_argument('--config', type=str, default=None,
help='path to the YAML config file.')
args = parser.parse_args()
if args.config is not None:
with open(args.config, 'r') as f:
config = yaml.safe_load(f)
args_dict = vars(args)
args_dict.update(config)
args_new = argparse.Namespace(**args_dict)
return args_new
return args
def bool_parser(arg):
"""Parses an argument to boolean."""
if isinstance(arg, bool):
return arg
if arg is None:
return False
if arg.lower() in ['1', 'true', 't', 'yes', 'y']:
return True
if arg.lower() in ['0', 'false', 'f', 'no', 'n']:
return False
raise ValueError(f'`{arg}` cannot be converted to boolean!')