I have been trying to convert the TensorFlow code into pytorch ,confused to what to use instead tf.keras.layer.Layer in pytorch
below is the code
“”"
Objective wrapper and utils to build a function to be optimized
“”"
import itertools
import tensorflow as tf
import numpy as np
from …commons import find_layer
from …types import Union, List, Callable, Tuple, Optional
from .losses import dot_cossim
class Objective:
“”"
Use to combine several sub-objectives into one.
Each sub-objective act on a layer, possibly on a neuron or a channel (in
that case we apply a mask on the layer values), or even multiple neurons (in
that case we have multiples masks). When two sub-objectives are added, we
optimize all their combinations.
e.g Objective 1 target the neurons 1 to 10 of the logits l1,...,l10
Objective 2 target a direction on the first layer d1
Objective 3 target each of the 5 channels on another layer c1,...,c5
The resulting Objective will have 10*5*1 combinations. The first input
will optimize l1+d1+c1 and the last one l10+d1+c5.
Parameters
----------
model
Model used for optimization.
layers
A list of the layers output for each sub-objectives.
masks
A list of masks that will be applied on the targeted layer for each
sub-objectives.
funcs
A list of loss functions for each sub-objectives.
multipliers
A list of multiplication factor for each sub-objectives
names
A list of name for each sub-objectives
"""
def __init__(self,
model: tf.keras.Model,
layers: List[tf.keras.layers.Layer],
masks: List[tf.Tensor],
funcs: List[Callable],
multipliers: List[float],
names: List[str]):
self.model = model
self.layers = layers
self.masks = masks
self.funcs = funcs
self.multipliers = multipliers
self.names = names
def __add__(self, term):
if not isinstance(term, Objective):
raise ValueError(f"{term} is not an objective.")
return Objective(
self.model,
layers=self.layers + term.layers,
masks=self.masks + term.masks,
funcs=self.funcs + term.funcs,
multipliers=self.multipliers + term.multipliers,
names=self.names + term.names
)
def __sub__(self, term):
if not isinstance(term, Objective):
raise ValueError(f"{term} is not an objective.")
term.multipliers = [-1.0 * m for m in term.multipliers]
return self + term
def __mul__(self, factor: float):
if not isinstance(factor, (int, float)):
raise ValueError(f"{factor} is not a number.")
self.multipliers = [m * factor for m in self.multipliers]
return self
def __rmul__(self, factor: float):
return self * factor
def compile(self) -> Tuple[tf.keras.Model, Callable, List[str], Tuple]:
"""
Compile all the sub-objectives into one and return the objects
for the optimisation process.
Returns
-------
model_reconfigured
Model with the outputs needed for the optimization.
objective_function
Function to call that compute the loss for the objectives.
names
Names of each objectives.
input_shape
Shape of the input, one sample for each optimization.
"""
# the number of inputs will be the number of combinations possible
# of the objectives, the mask are used to take into account
# these combinations
nb_sub_objectives = len(self.multipliers)
# re-arrange to match the different objectives with the model outputs
masks = np.array([np.array(m, dtype=object) for m in itertools.product(*self.masks)])
masks = [tf.cast(tf.stack(list(masks[:, i])), tf.float32) for i in
range(nb_sub_objectives)]
# the name of each combination is the concatenation of each objectives
names = np.array([' & '.join(names) for names in
itertools.product(*self.names)])
# one multiplier by sub-objective
multipliers = tf.constant(self.multipliers)
def objective_function(model_outputs):
loss = 0.0
for output_index in range(0, nb_sub_objectives):
outputs = model_outputs[output_index]
loss += self.funcs[output_index](
outputs, tf.cast(masks[output_index], outputs.dtype))
loss *= multipliers[output_index]
return loss
# the model outputs will be composed of the layers needed
model_reconfigured = tf.keras.Model(self.model.input, [*self.layers])
nb_combinations = masks[0].shape[0]
input_shape = (nb_combinations, *model_reconfigured.input.shape[1:])
return model_reconfigured, objective_function, names, input_shape
can anyone pls help me out with this??