I have been trying to convert the TensorFlow code into pytorch ,facing the problem ,what to use tf.keras.layer.Layer in pytorch ,Here below is the code ,could any pls help me out with this part?
"""
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
I have tried to convert the code torch and expecting the same out as TensorFlow code