import argparse import os import shutil import numpy as np import torch # from trixi.util import Config, GridSearch def check_attributes(object_, attributes): missing = [] for attr in attributes: if not hasattr(object_, attr): missing.append(attr) if len(missing) > 0: return False else: return True def set_seeds(seed, cuda=True): if not hasattr(seed, "__iter__"): seed = (seed, seed, seed) np.random.seed(seed[0]) torch.manual_seed(seed[1]) if cuda: torch.cuda.manual_seed_all(seed[2]) def make_onehot(array, labels=None, axis=1, newaxis=False): # get labels if necessary if labels is None: labels = np.unique(array) labels = list(map(lambda x: x.item(), labels)) # get target shape new_shape = list(array.shape) if newaxis: new_shape.insert(axis, len(labels)) else: new_shape[axis] = new_shape[axis] * len(labels) # make zero array if type(array) == np.ndarray: new_array = np.zeros(new_shape, dtype=array.dtype) elif torch.is_tensor(array): new_array = torch.zeros(new_shape, dtype=array.dtype, device=array.device) else: raise TypeError("Onehot conversion undefined for object of type {}".format(type(array))) # fill new array n_seg_channels = 1 if newaxis else array.shape[axis] for seg_channel in range(n_seg_channels): for l, label in enumerate(labels): new_slc = [slice(None), ] * len(new_shape) slc = [slice(None), ] * len(array.shape) new_slc[axis] = seg_channel * len(labels) + l if not newaxis: slc[axis] = seg_channel new_array[tuple(new_slc)] = array[tuple(slc)] == label return new_array def match_to(x, ref, keep_axes=(1,)): target_shape = list(ref.shape) for i in keep_axes: target_shape[i] = x.shape[i] target_shape = tuple(target_shape) if x.shape == target_shape: pass if x.dim() == 1: x = x.unsqueeze(0) if x.dim() == 2: while x.dim() < len(target_shape): x = x.unsqueeze(-1) x = x.expand(*target_shape) x = x.to(device=ref.device, dtype=ref.dtype) return x def make_slices(original_shape, patch_shape): working_shape = original_shape[-len(patch_shape):] splits = [] for i in range(len(working_shape)): splits.append([]) for j in range(working_shape[i] // patch_shape[i]): splits[i].append(slice(j*patch_shape[i], (j+1)*patch_shape[i])) rest = working_shape[i] % patch_shape[i] if rest > 0: splits[i].append(slice((j+1)*patch_shape[i], (j+1)*patch_shape[i] + rest)) # now we have all slices for the individual dimensions # we need their combinatorial combinations slices = list(itertools.product(*splits)) for i in range(len(slices)): slices[i] = [slice(None), ] * (len(original_shape) - len(patch_shape)) + list(slices[i]) return slices def coordinate_grid_samples(mean, std, factor_std=5, scale_std=1.): relative = np.linspace(-scale_std*factor_std, scale_std*factor_std, 2*factor_std+1) positions = np.array([mean + i * std for i in relative]).T axes = np.meshgrid(*positions) axes = map(lambda x: list(x.ravel()), axes) samples = list(zip(*axes)) samples = list(map(np.array, samples)) return samples def get_default_experiment_parser(): parser = argparse.ArgumentParser() parser.add_argument("base_dir", type=str, help="Working directory for experiment.") parser.add_argument("-c", "--config", type=str, default=None, help="Path to a config file.") parser.add_argument("-v", "--visdomlogger", action="store_true", help="Use visdomlogger.") parser.add_argument("-tx", "--tensorboardxlogger", type=str, default=None) parser.add_argument("-tl", "--telegramlogger", action="store_true") parser.add_argument("-dc", "--default_config", type=str, default="DEFAULTS", help="Select a default Config") parser.add_argument("-ad", "--automatic_description", action="store_true") parser.add_argument("-r", "--resume", type=str, default=None, help="Path to resume from") parser.add_argument("-irc", "--ignore_resume_config", action="store_true", help="Ignore Config in experiment we resume from.") parser.add_argument("-test", "--test", action="store_true", help="Run test instead of training") parser.add_argument("-g", "--grid", type=str, help="Path to a config for grid search") parser.add_argument("-s", "--skip_existing", action="store_true", help="Skip configs for which an experiment exists, only for grid search") parser.add_argument("-m", "--mods", type=str, nargs="+", default=None, help="Mods are Config stubs to update only relevant parts for a certain setup.") parser.add_argument("-ct", "--copy_test", action="store_true", help="Copy test files to original experiment.") return parser def run_experiment(experiment, configs, args, mods=None, **kwargs): # set a few defaults if "explogger_kwargs" not in kwargs: kwargs["explogger_kwargs"] = dict(folder_format="{experiment_name}_%Y%m%d-%H%M%S") if "explogger_freq" not in kwargs: kwargs["explogger_freq"] = 1 if "resume_save_types" not in kwargs: kwargs["resume_save_types"] = ("model", "simple", "th_vars", "results") config = Config(file_=args.config) if args.config is not None else Config() config.update_missing(configs[args.default_config].deepcopy()) if args.mods is not None and mods is not None: for mod in args.mods: config.update(mods[mod]) config = Config(config=config, update_from_argv=True) # GET EXISTING EXPERIMENTS TO BE ABLE TO SKIP CERTAIN CONFIGS if args.skip_existing: existing_configs = [] for exp in os.listdir(args.base_dir): try: existing_configs.append(Config(file_=os.path.join(args.base_dir, exp, "config", "config.json"))) except Exception as e: pass if args.grid is not None: grid = GridSearch().read(args.grid) else: grid = [{}] for combi in grid: config.update(combi) if args.skip_existing: skip_this = False for existing_config in existing_configs: if existing_config.contains(config): skip_this = True break if skip_this: continue if "backup_every" in config: kwargs["save_checkpoint_every_epoch"] = config["backup_every"] loggers = {} if args.visdomlogger: loggers["v"] = ("visdom", {}, 1) if args.tensorboardxlogger is not None: if args.tensorboardxlogger == "same": loggers["tx"] = ("tensorboard", {}, 1) else: loggers["tx"] = ("tensorboard", {"target_dir": args.tensorboardxlogger}, 1) if args.telegramlogger: kwargs["use_telegram"] = True if args.automatic_description: difference_to_default = Config.difference_config_static(config, configs["DEFAULTS"]).flat(keep_lists=True, max_split_size=0, flatten_int=True) description_str = "" for key, val in difference_to_default.items(): val = val[0] description_str = "{} = {}\n{}".format(key, val, description_str) config.description = description_str exp = experiment(config=config, base_dir=args.base_dir, resume=args.resume, ignore_resume_config=args.ignore_resume_config, loggers=loggers, **kwargs) trained = False if args.resume is None or args.test is False: exp.run() trained = True if args.test: exp.run_test(setup=not trained) if isinstance(args.resume, str) and exp.elog is not None and args.copy_test: for f in glob.glob(os.path.join(exp.elog.save_dir, "test*")): if os.path.isdir(f): shutil.copytree(f, os.path.join(args.resume, "save", os.path.basename(f))) else: shutil.copy(f, os.path.join(args.resume, "save"))