def gradcam(
features: Module, classifier: Module, input_image: Tensor, n_top_classes: int = 3
) -> Tensor:
"""Performs GradCAM on an input image.
For further explanations about GradCam, see https://arxiv.org/abs/1610.02391.
Args:
features: the spatial feature part of the model, before classifier
classifier: the classifier part of the model, after spatial features
input_image: image tensor of dimensions (c, h, w) or (1, c, h, w)
n_top_classes: the number of classes to calculate GradCAM for
Returns:
a GradCAM heatmap of dimensions (h, w)
"""
# Get selector for top k classes
input_image = _prepare_input(input_image)
class_selector = _top_k_selector(
Sequential(features.eval(), classifier.eval()), input_image, n_top_classes
)
# Apply spatial GradAM on classes
gradam = GradAM(classifier, class_selector, features, SpatialSplit())
result = gradam.visualize(input_image)
return _upscale(result, tuple(input_image.size()[2:]))
def remove(self, module: Module) -> None:
with torch.no_grad():
weight = self.compute_weight(module, do_power_iteration=False)
delattr(module, self.name)
delattr(module, self.name + "_u")
delattr(module, self.name + "_v")
delattr(module, self.name + "_orig")
module.register_parameter(self.name, torch.nn.Parameter(weight.detach()))
def create_optimizer(self,opts, model: Module):
optimopts = opts.optimizeropts
if opts.epocheropts.gpu:
device = torch.device("cpu")
model = model.to(device=device)
optim = globals()[optimopts.type](model.parameters(),
lr=optimopts.lr,
momentum=optimopts.momentum,
weight_decay=optimopts.weight_decay,
dampening=optimopts.dampening,
nesterov=optimopts.nestrov)
opts.optimizeropts.lr_sched = LambdaLR(optim, opts.optimizeropts.lr_sched_lambda, last_epoch=-1)
return optim
def load_pytorch_model(self,
model: Module,
identifier,
*tags,
checkpoint_number=None):
load_dir = build_checkpoint_path(AgentBase.CHECKPOINT_DIR_NAME,
AgentBase.PYTORCH_PLATFORM,
identifier, *tags)
if not checkpoint_number:
checkpoint_number = self.get_next_index(load_dir)
load_path = os.path.join(load_dir,
(AgentBase.FILENAME + AgentBase.SEPARATOR +
str(checkpoint_number)))
LOG.info("Loading pytorch model at " + load_path)
model.load_state_dict(torch.load(load_path))
def class_visualization(net: Module, class_index: int) -> Tensor:
"""Visualizes a class for a classification network.
Args:
net: the network to visualize for
class_index: the index of the class to visualize
"""
if class_index < 0:
raise ValueError(f"Invalid class: {class_index}")
img = PixelActivation(
net.eval(),
SplitSelector(NeuronSplit(), [class_index]),
opt_n=500,
iter_n=20,
init_size=50,
transform=RandomTransform(scale_fac=0)
+ BilateralTransform()
+ ResizeTransform(1.1),
regularization=[TVRegularization(5e1), WeightDecay(1e-9)],
).visualize()
return PixelActivation(
net,
SplitSelector(NeuronSplit(), [class_index]),
opt_n=100,
iter_n=int(50),
transform=RandomTransform() + BilateralTransform(),
regularization=[TVRegularization(), WeightDecay()],
).visualize(img)
def validate(self, model: Module, data_loader: DataLoader, use_cuda: bool,
criterion: Module) -> float:
"""Performs one epoch of validating of the model and returns the
obtained validation loss.
:param model: model to be validated
:param data_loader: validation data loader
:param use_cuda: a flag whether CUDA can be used
:param criterion: loss function
:return: validation loss
"""
valid_loss = 0.0
model.eval()
for batch_idx, (data, target) in enumerate(tqdm(data_loader)):
data = self.move_to_gpu(data, use_cuda)
output = model.forward(*data)
loss = criterion(*output)
valid_loss += ((1 / (batch_idx + 1)) * (loss.data - valid_loss))
return valid_loss
def save_pytorch_model(self, model: Module, identifier, *tags):
save_dir = build_checkpoint_path(AgentBase.CHECKPOINT_DIR_NAME,
AgentBase.PYTORCH_PLATFORM,
identifier, *tags)
self.init_dir(save_dir)
save_path = os.path.join(save_dir,
(AgentBase.FILENAME + AgentBase.SEPARATOR +
str(self.get_next_index(save_dir))))
LOG.info("Saving pytorch model at " + save_path)
torch.save(model.state_dict(), save_path)
def sgd(
model: Module,
lr: float = 0.01,
momentum: float = 0.9,
weight_decay: float = 5e-4,
) -> OptimizerSchedulerBundle:
optimizer = SGD(model.parameters(),
lr=lr,
momentum=momentum,
weight_decay=weight_decay)
return OptimizerSchedulerBundle(optimizer=optimizer)
def lrs(model: Module,
lr: float = 0.1,
momentum: float = 0.9,
weight_decay: float = 5e-4,
step_size: int = 30) -> OptimizerSchedulerBundle:
optimizer = SGD(model.parameters(),
lr=lr,
momentum=momentum,
weight_decay=weight_decay)
scheduler = StepLR(optimizer, step_size=step_size)
return OptimizerSchedulerBundle(optimizer=optimizer, scheduler=scheduler)
def train(self, model: Module, data_loader: DataLoader, use_cuda: bool,
criterion: Module, optimizer: Optimizer) -> float:
"""Performs one epoch of training of the model and returns the obtained
training loss.
:param model: model to be trained
:param data_loader: training data loader
:param use_cuda: a flag whether CUDA can be used
:param criterion: loss function
:param optimizer: optimizer
:return: training loss
"""
train_loss = 0.0
model.train()
for batch_idx, (data, target) in enumerate(tqdm(data_loader)):
data = self.move_to_gpu(data, use_cuda)
optimizer.zero_grad()
output = model(*data)
loss = criterion(*output)
loss.backward()
optimizer.step()
train_loss += ((1 / (batch_idx + 1)) * (loss.data - train_loss))
return train_loss
def _top_k_selector(net: Module, img: Tensor, k: int) -> NeuronSelector:
"""Creates a top-k classes selector.
Args:
net: network
img: prepared input image
k: number of classes
Returns:
neuron selector for the top k classes
"""
net = net.to(midnite.get_device())
out = net(img).squeeze(0)
mask = _top_k_mask(out, k)
return SimpleSelector(mask)
def occlusion(net: Module, input_image: Tensor, n_top_classes: int = 3) -> Tensor:
"""Creates a attribution heatmap by occluding parts of the input image.
Args:
net: the network to visualize attribution for
input_image: image tensor of dimensions (c, h, w) or (1, c, h, w)
n_top_classes: the number of classes to account for
Returns:
a occlusion heatmap of dimensions (h, w)
"""
input_image = _prepare_input(input_image)
class_selector = _top_k_selector(net.eval(), input_image, n_top_classes)
# Apply occlusion
occlusion_ = Occlusion(net, class_selector, SpatialSplit(), [1, 10, 10], [1, 5, 5])
result = occlusion_.visualize(input_image)
return _upscale(result, tuple(input_image.size()[2:]))
def apply(
module: Module, name: str, coeff: float, n_power_iterations: int, dim: int, eps: float,
) -> "SpectralNorm":
for k, hook in module._forward_pre_hooks.items():
if isinstance(hook, SpectralNorm) and hook.name == name:
raise RuntimeError("Cannot register two spectral_norm hooks on " "the same parameter {}".format(name))
fn = SpectralNorm(coeff, name, n_power_iterations, dim, eps)
weight = module._parameters[name]
with torch.no_grad():
weight_mat = fn.reshape_weight_to_matrix(weight)
h, w = weight_mat.size()
# randomly initialize `u` and `v`
u = normalize(weight.new_empty(h).normal_(0, 1), dim=0, eps=fn.eps)
v = normalize(weight.new_empty(w).normal_(0, 1), dim=0, eps=fn.eps)
delattr(module, fn.name)
module.register_parameter(fn.name + "_orig", weight)
# We still need to assign weight back as fn.name because all sorts of
# things may assume that it exists, e.g., when initializing weights.
# However, we can't directly assign as it could be an nn.Parameter and
# gets added as a parameter. Instead, we register weight.data as a plain
# attribute.
setattr(module, fn.name, weight.data)
module.register_buffer(fn.name + "_u", u)
module.register_buffer(fn.name + "_v", v)
module.register_buffer(fn.name + "_sigma", torch.ones(1).to(weight.device))
module.register_forward_pre_hook(fn)
module._register_state_dict_hook(SpectralNormStateDictHook(fn))
module._register_load_state_dict_pre_hook(SpectralNormLoadStateDictPreHook(fn))
return fn
def radam(model: Module) -> OptimizerSchedulerBundle:
optimizer = RAdam(model.parameters())
return OptimizerSchedulerBundle(optimizer=optimizer)