def __init__(self, optims, initial, decay_fac, decay_interval_itr=None, decay_interval_epoch=None, epoch_len=None, warm_restart=None, warm_restart_schedule=None): super(ExponentialDecayLRSchedule, self).__init__(optims) assert_exc((decay_interval_itr is not None) ^ (decay_interval_epoch is not None), 'Need either iter or epoch') if decay_interval_epoch: assert epoch_len is not None decay_interval_itr = int(decay_interval_epoch * epoch_len) if warm_restart: warm_restart = int(warm_restart * epoch_len) self.initial = initial self.decay_fac = decay_fac self.decay_every_itr = decay_interval_itr self.warm_restart_itr = warm_restart self.warm_restart_schedule = warm_restart_schedule self.last_warm_restart = 0
def _enable(self, prefix, global_step): """ Enable logging of prefix """ assert_exc(isinstance(prefix, str), 'prefix must be str, got {}'.format(prefix)) assert_exc(prefix[-1] != '/') self.enabled_prefix = prefix self.global_step = global_step
def get_ckpt_for_itr(self, itr, before_time=None): """ Gets ckpt_itrc where itrc <= itr, i.e., the latest ckpt before `itr`. If `before_time` is given and itr == -1, the latest ckpt before `before_time` is used Special values: itr == -1 -> newest ckpt """ # sorted list of (itr, ckpt_p) ckpts = list(self.itr_ckpt()) if before_time is not None: ckpts_before_time = [(i, p) for i, p in ckpts if os.path.getmtime(p) <= before_time] print( f'*** Ignoring {len(ckpts) - len(ckpts_before_time)} ckpts after {before_time}' ) ckpts = ckpts_before_time assert_exc( len(ckpts) > 0, 'No ckpts found in {}'.format(self._out_dir), CkeckpointLoadingException) if itr == -1: return ckpts[-1] first_itrc, _ = ckpts[0] assert_exc(first_itrc <= itr, 'Earliest ckpt {} is after {}'.format(first_itrc, itr), CkeckpointLoadingException) for itrc, ckpt_p in reversed(ckpts): if itrc <= itr: return itrc, ckpt_p raise ValueError('Unexpected, {}, {}'.format(itr, ckpts))
def _parse_cos_spec(): lrmax, lrmin, T = s.split(SPEC_SEP) kind, T = T[0], T[1:] assert_exc(kind in ('i', 'e'), 'Invalid spec: {}'.format(s)) T_itr = int(T) if kind == 'i' else None T_epoch = float(T) if kind == 'e' else None return CosineDecayLRSchedule(optims, float(lrmax), float(lrmin), T_itr, T_epoch, epoch_len)
def add(self, v): v = pe.tensor_to_np(v) num_values = np.prod(v.shape) if self._buffer is None: print(f'Creating {v.dtype} buffer for {self._name}: {self._buffer_size}x{num_values}') self._buffer = np.zeros((self._buffer_size, num_values), dtype=v.dtype) assert_exc(self._buffer.shape[1] == num_values, (self._buffer.shape, v.shape, num_values), BufferSizeMismatch) self._buffer[self._idx, :] = v.flatten() self._idx = (self._idx + 1) % self._buffer_size self._filled_idx = min(self._filled_idx + 1, self._buffer_size)
def __getitem__(self, items): """ Allows usage as lr, L, num_layers = config['lr', 'L', 'num_layer'] """ if not isinstance(items, tuple): items = [items] for item in items: assert_exc(item in self.__dict__, 'Invalid parameter: {}'.format(item), AttributeError) yield self.__dict__[item]
def get_ckpt_for_itr(self, itr): """ Gets ckpt_itrc where itrc <= itr, i.e., the latest ckpt before `itr`. Special values: itr == -1 -> newest ckpt """ ckpts = list(self.itr_ckpt()) assert_exc(len(ckpts) > 0, 'No ckpts found in {}'.format(self._out_dir)) if itr == -1: return ckpts[-1] first_itrc, _ = ckpts[0] assert_exc(first_itrc <= itr, 'Earliest ckpt {} is after {}'.format(first_itrc, itr)) for itrc, ckpt_p in reversed(ckpts): if itrc <= itr: return itrc, ckpt_p raise ValueError('Unexpected, {}, {}'.format(itr, ckpts))
def decode(self, pin, png_out_p): """ Decode L3C-encoded file at `pin` to a PNG at `png_out_p`. """ pout_dir = os.path.dirname(os.path.abspath(png_out_p)) assert_exc(os.path.isdir(pout_dir), f'png_out_p directory ({pout_dir}) does not exists!', DecodeError) assert_exc(png_out_p.endswith('.png'), f'png_out_p must end in .png, got {png_out_p}', DecodeError) decoded = self.bc.decode(pin) self._write_img(decoded, png_out_p) print(f'---\nDecoded: {png_out_p}')
def encode(self, img_p, pout, overwrite=False): pout_dir = os.path.dirname(os.path.abspath(pout)) assert_exc(os.path.isdir(pout_dir), f'pout directory ({pout_dir}) does not exists!', EncodeError) if overwrite and os.path.isfile(pout): print(f'Removing {pout}...') os.remove(pout) assert_exc(not os.path.isfile(pout), f'{pout} exists. Consider --overwrite', EncodeError) img = self._read_img(img_p) img = img.to(pe.DEVICE) self.bc.encode(img, pout=pout) print('---\nSaved:', pout)
def __init__(self, root_dir_or_img, max_imgs=None, skip_hidden=False, append_id=None): """ :param root_dir_or_img: Either a directory with images or the path of a single image. :param max_imgs: If given, subsample deterministically to only contain max_imgs :param skip_hidden: If given, skip images starting with '.' :param append_id: If given, append `append_id` to self.id :raises ValueError if root_dir is not a directory or does not contain images """ self.root_dir_or_img = root_dir_or_img if os.path.isdir(root_dir_or_img): root_dir = root_dir_or_img self.name = os.path.basename(root_dir.rstrip('/')) self.ps = sorted(p for p in os_ext.listdir_paths(root_dir) if has_image_ext(p)) if skip_hidden: self.ps = self._filter_hidden(self.ps) if max_imgs and max_imgs < len(self.ps): print('Subsampling to use {} imgs of {}...'.format( max_imgs, self.name)) idxs = np.linspace(0, len(self.ps) - 1, max_imgs, dtype=np.int) self.ps = np.array(self.ps)[idxs].tolist() assert len(self.ps) == max_imgs assert_exc( len(self.ps) > 0, 'No images found in {}'.format(root_dir), ValueError) self.id = '{}_{}'.format(self.name, len(self.ps)) self._str = 'Testset({}): in {}, {} images'.format( self.name, root_dir, len(self.ps)) else: img = root_dir_or_img assert_exc(os.path.isfile(img), 'Does not exist: {}'.format(img), FileNotFoundError) self.name = os.path.basename(img) self.ps = [img] self.id = img self._str = 'Testset([{}]): 1 image'.format(self.name) if append_id: self.id += append_id
def interpolator(measures_per_image_iter, grid, interp_mode='linear'): accumulated_values = np.zeros_like(grid, np.float64) # Count values per bin N = np.zeros_like(grid, np.int64) num_imgs = 0 num_errors = 0 for img_description, (bpps, values) in measures_per_image_iter: assert_exc( len(bpps) >= 2, 'Missing values for {}'.format(img_description), OtherCodecsReadException) assert_exc(bpps[0] >= bpps[-1], f'First bpp < last: {bpps[0]} < {bpps[-1]}', OtherCodecsReadException) num_imgs += 1 # create interpolation function try: fq = scipy.interpolate.interp1d(bpps, values, interp_mode) except ValueError as e: raise OtherCodecsReadException('ValueError while creating fq:', e) for i, bpp in enumerate(grid): try: accumulated_values[i] += fq(bpp) N[i] += 1 except ValueError as e: num_errors += 1 continue try: grid, values = ft.unzip( (bpp, v / n) for bpp, v, n in zip(grid, accumulated_values, N) if n > _REQUIRED_BINS * num_imgs) except ValueError as e: raise ValueError( grid, accumulated_values, N, list( ft.unzip((bpp, v / n) for bpp, v, n in zip(grid, accumulated_values, N) if n > _REQUIRED_BINS * num_imgs))) return grid, values
def to_image(t): """ :param t: tensor or np.ndarray, may be of shape NCHW / CHW with C=1 or 3 / HW, dtype float32 or uint8. If float32: must be in [0, 1] :return: HW3 uint8 np.ndarray """ if not isinstance(t, np.ndarray): t = pe.tensor_to_np(t) # - t is numpy array if t.ndim == 4: # - t has batch dimension, only use first t = t[0, ...] elif t.ndim == 2: t = np.expand_dims(t, 0) # Now 1HW assert_exc(t.ndim == 3, 'Invalid shape: {}'.format(t.shape)) # - t is 3 dimensional CHW numpy array if t.dtype != np.uint8: assert_exc(t.dtype == np.float32, 'Expected either uint8 or float32, got {}'.format(t.dtype)) _check_range(t, 0, 1) t = (t * 255.).astype(np.uint8) # - t is uint8 numpy array num_channels = t.shape[0] if num_channels == 3: t = np.transpose(t, (1, 2, 0)) elif num_channels == 1: t = np.stack([t[0, :, :] for _ in range(3)], -1) else: raise ValueError('Expected CHW, got {}'.format(t.shape)) assert_exc(t.ndim == 3 and t.shape[2] == 3, str(t.shape)) # - t is uint8 numpy array of shape HW3 return t
def create_curves_for_images(root_dir, out_dir, grid, mode): times = [] # make sure we exclude here! otherwise, task_array.job_enumerate sees different number of files in every job!! all_img_ps = _get_image_paths(root_dir, exclude={'tmp'}) assert_exc(len(all_img_ps) > 0, 'No images found', ValueError) non_pngs = [p for p in all_img_ps if not p.endswith('.png')] assert_exc( len(non_pngs) == 0, f'Only .pngs are supported by this code! Found {len(non_pngs)} others.' ) measure_over_interval = { 'bpg': bpg_measure_over_interval, 'balle': balle_measure_over_interval, # 'bpgslow': bpg_measure_over_interval_slow, 'jp2k': jp2k_measure_over_interval, 'jp': jp_measure_over_interval, 'webp': webp_measure_over_interval }[mode] for i, img_p in task_array.job_enumerate(all_img_ps): print('>>>', task_array.TASK_ID, 'compresses', os.path.basename(img_p)) img_name = os.path.splitext(os.path.basename(img_p))[0] s = time.time() mf = measures_file_p(out_dir, img_name) if complete_measures_file_exists(mf, num_ops=len(grid)): print(f'Found output for {img_name}, skipping...') continue # need to create measures file with open(mf, 'w+') as f: measure_over_interval(img_p, f, grid) times.append(time.time() - s) avg_time = np.mean(times[-15:]) print('Time left: {:.2f}min'.format(avg_time * (len(all_img_ps) - i) / 60))
def _parse_exp_spec(s, optims, initial_lr, epoch_len): if s.count(SPEC_SEP) > 2: fac, interval, warm, warm_start, warm_fac, warm_interval = s.split( SPEC_SEP) assert warm == 'warm' warm_start = int(warm_start) warm_schedule = _parse_exp_spec( SPEC_SEP.join([warm_fac, warm_interval]), optims, initial_lr, epoch_len) else: fac, interval = s.split(SPEC_SEP) warm_start, warm_schedule = None, None kind, interval = interval[0], interval[1:] assert_exc(kind in ('i', 'e'), 'Invalid spec: {}'.format(s)) decay_interval_itr = int(interval) if kind == 'i' else None decay_interval_epoch = float(interval) if kind == 'e' else None return ExponentialDecayLRSchedule(optims, initial_lr, float(fac), decay_interval_itr, decay_interval_epoch, epoch_len, warm_restart=warm_start, warm_restart_schedule=warm_schedule)
def get_experiment_dir(log_dir, experiment_spec): """ experiment_spec: if is a logdate, find correct full path in log_dir, otherwise assume logdir/experiment_spec exists :return experiment dir, no slash at the end. containing /ckpts """ if logdir_helpers.is_log_date(experiment_spec): # assume that log_dir/restore* matches assert_exc(log_dir is not None, 'Can only infer experiment_dir from log_date if log_dir is not None') restore_dir_glob = os.path.join(log_dir, experiment_spec + '*') restore_dir_possible = glob.glob(restore_dir_glob) assert_exc(len(restore_dir_possible) == 1, 'Expected one match for {}, got {}'.format( restore_dir_glob, restore_dir_possible)) experiment_spec = restore_dir_possible[0] else: experiment_spec = os.path.join(log_dir, experiment_spec) experiment_dir = experiment_spec.rstrip(os.path.sep) assert_exc(os.path.isdir(experiment_dir), 'Invalid experiment_dir: {}'.format(experiment_dir)) return experiment_dir
def get_ckpts_dir(experiment_dir, ensure_exists=True): ckpts_p = os.path.join(experiment_dir, CKPTS_DIR_NAME) if ensure_exists: assert_exc(os.path.isdir(ckpts_p), 'Not found: {}'.format(ckpts_p)) return ckpts_p
def _assert_dir_exists_in_cwd(p): assert_exc(os.path.isdir(p), '{} not in {}'.format(p, os.getcwd()))
def filter_filenames(self, filter_filenames): filename = lambda p: os.path.splitext(os.path.basename(p))[0] self.ps = [p for p in self.ps if filename(p) in filter_filenames] assert_exc( len(self.ps) > 0, 'No files after filtering for {}'.format(filter_filenames))
def _check_range(a, lo, hi): a_lo, a_hi = np.min(a), np.max(a) assert_exc(a_lo >= lo and a_hi <= hi, 'Invalid range: [{}, {}]. Expected: [{}, {}]'.format(a_lo, a_hi, lo, hi))