File size: 8,409 Bytes
acc22af |
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 |
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")) |