DL-Art-School/codes/models/experiments/experiments.py
2020-09-19 10:05:25 -06:00

83 lines
3.0 KiB
Python

import torch
def get_experiment_for_name(name):
return Experiment()
# Experiments are ways to add hooks into the ExtensibleTrainer training process with the intent of reporting the
# inner workings of the process in a custom manner that is unsuitable for addition elsewhere.
class Experiment:
def before_step(self, opt, step_name, env, nets_to_train, pre_state):
pass
def before_optimize(self, state):
pass
def after_optimize(self, state):
pass
def get_log_data(self):
pass
class ModelParameterDepthTrackerMetrics(Experiment):
# Subclasses should implement these two methods:
def get_network_and_step_names(self):
# Subclasses should return the network being debugged and the step name it is trained in. return: (net, stepname)
pass
def get_layers_to_debug(self, env, net, state):
# Subclasses should populate self.layers with a list of per-layer nn.Modules here.
pass
def before_step(self, opt, step_name, env, nets_to_train, pre_state):
self.net, step = self.get_network_and_step_names()
self.activate = self.net in nets_to_train and step == step_name and self.step_num % opt['logger']['print_freq'] == 0
if self.activate:
layers = self.get_layers_to_debug(env, env['networks'][self.net], pre_state)
self.params = []
for l in layers:
lparams = []
for k, v in env['networks'][self.net].named_parameters(): # can optimize for a part of the model
if v.requires_grad:
lparams.append(v)
self.params.append(lparams)
def before_optimize(self, state):
self.cached_params = []
for l in self.params:
lparams = []
for p in l:
lparams.append(p.value().cpu())
self.cached_params.append(lparams)
def after_optimize(self, state):
# Compute the abs mean difference across the params.
self.layer_means = []
for l, lc in zip(self.params, self.cached_params):
sum = torch.tensor(0)
for p, pc in zip(l, lc):
sum += torch.abs(pc - p.value().cpu())
self.layer_means.append(sum / len(l))
def get_log_data(self):
return {'%s_layer_update_means_histogram' % (self.net,): self.layer_means}
class DiscriminatorParameterTracker(ModelParameterDepthTrackerMetrics):
def get_network_and_step_names(self):
return "feature_discriminator", "feature_discriminator"
def get_layers_to_debug(self, env, net, state):
return [net.ref_head.conv0_0,
net.ref_head.conv0_1,
net.ref_head.conv1_0,
net.ref_head.conv1_1,
net.ref_head.conv2_0,
net.ref_head.conv2_1,
net.ref_head.conv3_0,
net.ref_head.conv3_1,
net.ref_head.conv4_0,
net.ref_head.conv4_1,
net.linear1,
net.output_linears]