def find_minimum(model: models.ModelWrapper,
optim_config: config.OptimConfig) -> dict:
optimiser = optim_config.algorithm_type(
model.parameters(),
**optim_config.algorithm_args) # type: optim.Optimizer
# Scheduler
if optim_config.scheduler_type is not None:
scheduler = optim_config.scheduler_type(optimiser,
**optim_config.scheduler_args)
else:
scheduler = None
# Initialise
model.initialise_randomly()
if optim_config.eval_config is not None:
model.adapt_to_config(optim_config.eval_config)
# Optimise
for _ in helper.pbar(range(optim_config.nsteps), "Find mimimum"):
model.apply(gradient=True)
if scheduler is not None:
scheduler.step()
optimiser.step()
# todo tensorboard logging or similar
result = {
"coords": model.get_coords().to("cpu"),
}
# Analyse
analysis = model.analyse()
logger.debug(f"Found minimum: {analysis}.")
result.update(analysis)
return result
def auto_neb(m1,
m2,
graph: MultiGraph,
model: models.ModelWrapper,
config: config.AutoNEBConfig,
callback: callable = None):
# Continue existing cycles or start from scratch
if m2 in graph[m1]:
existing_edges = graph[m1][m2]
previous_cycle_idx = max(existing_edges)
connection_data = graph[m1][m2][previous_cycle_idx]
start_cycle_idx = previous_cycle_idx + 1
else:
connection_data = {
"path_coords":
torch.cat([graph.nodes[m]["coords"].view(1, -1)
for m in (m1, m2)]),
"target_distances":
torch.ones(1)
}
start_cycle_idx = 1
assert start_cycle_idx <= config.cycle_count
# Run NEB and add to graph
for cycle_idx in helper.pbar(
range(start_cycle_idx, config.cycle_count + 1), "AutoNEB"):
cycle_config = config.neb_configs[cycle_idx - 1]
connection_data = neb(connection_data, model, cycle_config)
graph.add_edge(m1,
m2,
key=cycle_idx,
**helper.move_to(connection_data, "cpu"))
if callback is not None:
callback()
def landscape_exploration(graph: MultiGraph,
model: models.ModelWrapper,
lex_config: config.LandscapeExplorationConfig,
callback: callable = None):
weight_key = lex_config.weight_key
try:
with helper.pbar(desc="Landscape Exploration") as bar:
while True:
# Suggest new pair based on current graph
m1, m2 = suggest_pair(graph, lex_config)
if m1 is None or m2 is None:
break
auto_neb(m1,
m2,
graph,
model,
lex_config.auto_neb_config,
callback=callback)
bar.update()
# Analyse new saddle
simple_graph = to_simple_graph(graph, weight_key)
best_saddle = simple_graph[m1][m2][weight_key]
in_mst_str = "included" if minimum_spanning_tree(
simple_graph, weight_key) else "not included"
logger.info(
f"Saddle loss between {m1} and {m2} is {best_saddle}, {in_mst_str} in MST."
)
finally:
simple_graph = to_simple_graph(graph, weight_key)
mst_graph = minimum_spanning_tree(simple_graph, weight_key)
mean_saddle_loss = sum(
simple_graph.get_edge_data(*edge)[weight_key]
for edge in mst_graph.edges) / len(mst_graph.edges)
logger.info(f"Average loss in MST: {mean_saddle_loss}.")
def main():
parser = ArgumentParser()
parser.add_argument("project_directory", nargs=1)
parser.add_argument("config_file", nargs=1)
parser.add_argument("--no-backup", default=False, action="store_true")
args = parser.parse_args()
project_directory = args.project_directory[0]
config_file = args.config_file[0]
graph_path, project_config_path = setup_project(config_file,
project_directory)
model, minima_count, min_config, lex_config = read_config_file(
project_config_path)
# Setup Logger
root_logger = getLogger()
root_logger.setLevel(INFO)
root_logger.addHandler(StreamHandler(sys.stdout))
root_logger.addHandler(
FileHandler(join(project_directory, "exploration.log")))
# === Create/load graph ===
if isfile(graph_path):
if not args.no_backup:
root_logger.info("Copying current 'graph.p' to backup file.")
copyfile(
graph_path,
graph_path.replace(".p", f"_bak{strftime('%Y%m%d-%H%M')}.p"))
else:
root_logger.info(
"Not creating a backup of 'graph.p' because of user request.")
graph = repair_graph(load_pickle_graph(graph_path), model)
else:
graph = MultiGraph()
# Call this after every optmisiation
def save_callback():
store_pickle_graph(graph, graph_path)
# === Ensure the specified number of minima ===
for _ in pbar(range(len(graph.nodes), minima_count), "Finding minima"):
minimum_data = find_minimum(model, min_config)
graph.add_node(
max(graph.nodes) + 1 if len(graph.nodes) > 0 else 1,
**move_to(minimum_data, "cpu"))
save_callback()
# === Connect minima ===
landscape_exploration(graph, model, lex_config, callback=save_callback)
def neb(previous_cycle_data, model: models.ModelWrapper,
neb_config: config.NEBConfig) -> dict:
# Initialise chain by inserting pivots
start_path, target_distances = neb_config.insert_method(
previous_cycle_data, **neb_config.insert_args)
# Model
neb_mod = neb_model.NEB(model, start_path, target_distances)
neb_mod.adapt_to_config(neb_config)
# Load optimiser
optim_config = neb_config.optim_config
# HACK: Optimisers only like parameters registered to autograd -> proper solution would keep several model instances as path and nudge their gradients after backward.
neb_mod.path_coords.requires_grad_(True)
optimiser = optim_config.algorithm_type(
neb_mod.parameters(),
**optim_config.algorithm_args) # type: optim.Optimizer
# HACK END: We don't want autograd to mingle with our computations
neb_mod.path_coords.requires_grad_(False)
if "weight_decay" in optimiser.defaults:
assert optimiser.defaults[
"weight_decay"] == 0, "NEB is not compatible with weight decay on the optimiser. Set weight decay on NEB instead."
# Scheduler
if optim_config.scheduler_type is not None:
scheduler = optim_config.scheduler_type(optimiser,
**optim_config.scheduler_args)
else:
scheduler = None
# Optimise
for _ in helper.pbar(range(optim_config.nsteps), "NEB"):
neb_mod.apply(gradient=True)
if scheduler is not None:
scheduler.step()
optimiser.step()
result = {
"path_coords": neb_mod.path_coords.clone().to("cpu"),
"target_distances": target_distances.to("cpu")
}
# Analyse
analysis = neb_mod.analyse(neb_config.subsample_pivot_count)
saddle_analysis = {
key: value
for key, value in analysis.items() if "saddle_" in key
}
logger.debug(f"Found saddle: {saddle_analysis}.")
result.update(analysis)
return result
def generate_dataset(self,
dataset,
number_of_classes,
valid_set_size=50,
adversarial_set_size=50):
final_set = []
adversarial_set = []
base_model = self.base_model
# Setup for adversarial samples
base_model.eval()
fool_model = foolbox.models.PyTorchModel(
base_model, (0, 1),
number_of_classes,
cuda=self.input_holder.device.type == "cuda")
criterion = foolbox.criteria.TargetClassProbability(self.label, p=0.99)
attack = foolbox.attacks.LBFGSAttack(fool_model, criterion)
for i in pbar(range(valid_set_size + adversarial_set_size),
desc="Generating dataset"):
data, target = dataset[i]
if target == self.label:
if len(final_set) < valid_set_size:
# Analyse the current image
self.input_holder.data[:] = data.view(
self.input_holder.shape)
user_data = self.analyse()
final_set.append((data, user_data))
else:
if len(adversarial_set) < adversarial_set_size:
# Generate an adversarial
adversarial = attack(data.numpy(), label=target)
# Skip failed attempts
if np.count_nonzero(adversarial == data.numpy()) == reduce(
mul, data.shape):
print("Failed to generate an adversarial example")
continue
data = from_numpy(adversarial).view(
self.input_holder.shape)
self.input_holder.data[:] = data
user_data = self.analyse()
user_data.update({
"adversarial": True,
"original_class": target
})
adversarial_set.append((data, user_data))
return adversarial_set + final_set
def analyse(self, sub_pivot_count=9):
# Collect stats here
analysis = {}
dense_pivot_count = (self.path_coords.shape[0] -
1) * (sub_pivot_count + 1) + 1
alphas = linspace(0, 1,
sub_pivot_count + 2)[:-1].to(self.path_coords.device)
for i in helpers.pbar(range(dense_pivot_count), "Saddle analysis"):
base_pivot = i // (sub_pivot_count + 1)
sub_pivot = i % (sub_pivot_count + 1)
if sub_pivot == 0:
# Coords of pivot
coords = self.path_coords[base_pivot]
else:
# Or interpolation between pivots
alpha = alphas[sub_pivot]
coords = self.path_coords[base_pivot] * (
1 - alpha) + self.path_coords[base_pivot + 1] * alpha
# Retrieve values from model analysis
self.model.set_coords_no_grad(coords)
point_analysis = self.model.analyse()
for key, value in point_analysis.items():
dense_key = "dense_" + key
if not isinstance(value, Tensor):
value = Tensor([value]).squeeze()
if dense_key not in analysis:
analysis[dense_key] = value.new(dense_pivot_count,
*value.shape)
analysis[dense_key][i] = value
# Compute saddle values
for key, value in list(analysis.items()):
if len(value.shape) == 1 or value.shape[1] == 1:
analysis[key.replace("dense_", "saddle_")] = value.max().item()
else:
print(key)
# Compute lengths
end_to_end_distance = (self.path_coords[-1] -
self.path_coords[0]).norm(2)
analysis["lengths"] = (self.path_coords[1:] - self.path_coords[:-1]
).norm(2, 1) / end_to_end_distance
analysis["length"] = end_to_end_distance
return analysis
def iterate_densely(self, sub_pivot_count=9):
dense_pivot_count = (self.path_coords.shape[0] -
1) * (sub_pivot_count + 1) + 1
alphas = linspace(0, 1,
sub_pivot_count + 2)[:-1].to(self.path_coords.device)
for i in helpers.pbar(range(dense_pivot_count), "Saddle analysis"):
base_pivot = i // (sub_pivot_count + 1)
sub_pivot = i % (sub_pivot_count + 1)
if sub_pivot == 0:
# Coords of pivot
coords = self.path_coords[base_pivot]
else:
# Or interpolation between pivots
alpha = alphas[sub_pivot]
coords = self.path_coords[base_pivot] * (
1 - alpha) + self.path_coords[base_pivot + 1] * alpha
# Retrieve values from model analysis
self.model.set_coords_no_grad(coords)
yield i