We Keep Coding

How to return extra loss from module forward function in PyTorch?

I made a module, that needs an extra loss term, e.g.
class MyModule:
def forward(self, x):
out = f(x)
extra_loss = loss_f(self.parameters(), x)
return out, extra_loss
I can't figure out how to make this module embeddable, for example, into a Sequential model: any regular module like Linear put after this one will fail because extra_loss causes the input to Linear to be a tuple, which Linear does not support.
So what I am looking for is extracting that extra loss after running the model forward
my_module = MyModule()
model = Sequential(
my_module,
Linear(my_module_outputs, 1)
)
output = model(x)
my_module_loss = ????
loss = mse(label, output) + my_module_loss
Does module composability support this scenario?

IMHO, hooks here is overreaction. Provided extra_loss is additive, we can use global variable like this:
class MyModule:
extra_loss =0
def forward(self, x):
out = f(x)
MyModule.extra_loss += loss_f(self.parameters(), x)
return out
output = model(x)
loss = mse(label, output) + MyModule.extra_loss
MyModule.extra_loss =0

You can register a hook in this case. A hook can be registered on a Tensor or a nn.Module. A hook is a function that is executed when the either forward or backward is called. In this case, we want to attach a forward hook without deattaching itself from the graph so that backward pass can happen.
import torch.nn as nn
act_out = {}
def get_hook(name):
def hook(m, input, output):
act_out[name] = output
return hook
class MyModule(torch.nn.Module):
def __init__(self, input, out, device=None):
super().__init__()
self.model = nn.Linear(input,out)
def forward(self,x):
return self.model(x), torch.sum(x) #our extra loss
class MyModule1(torch.nn.Module):
def __init__(self, input, out, device=None):
super().__init__()
self.model = nn.Linear(input,out)
def forward(self, pair):
x, loss = pair
return self.model(x)
model = nn.Sequential(
MyModule(5,10),
MyModule1(10,1)
)
for name, module in model.named_children():
print(name, module)
if name == '0':
module.register_forward_hook(get_hook(name))
x = torch.tensor([1,2,3,4,5]).float()
out = model(x)
print(act_out)
loss = myanotherloss(out)+act_out['0'][1] # this is the extra loss
# further processing
Note: I am using name == '0' because this is the only module where I want to attach the hook.
Note: Another notable point is nn.Sequential doesn't allow multiple inputs. In this case, it is simply considered as a tuple and then from that tuple we are using the loss and the input.

How be sure on_epoch_end triggers

I made a simple HDF5 reader class to avoid loading entire dataset in memory. I used sequence class to do so, but I'm not sure if on_epoch_end() function will trigger correctly.
I put a single print inside it, but it never appear! So I think there is something wrong in my code:
class HDF5Generator(tf.keras.utils.Sequence):
def __init__(self, hdf5_file, shuffle=True):
print("GENERATED")
self.hdf5 = h5py.File(hdf5_file, 'r')
self.shuffle = shuffle
self.indices = list(range(0, len(self.hdf5["samples"])))
random.Random().shuffle(self.indices)
def __len__(self):
return len(self.hdf5["samples"])
def __getitem__(self, idx):
return self.hdf5["samples"][self.indices[idx]], self.hdf5["labels"][self.indices[idx]]
def on_epoch_end(self):
print("RE-SHUFFLE")
random.Random().shuffle(self.indices)
Here how I use it:
d = tf.data.Dataset.from_generator(HD5FGenerator, args=[dataset], output_signature=(...))
d = d.batch(batch_size).prefetch(tf.data.AUTOTUNE).cache()
...
model.fit(d, epochs=epochs)
In console appear the epoch counter, the progress bar, the string "GENERATED" but never "RE-SHUFFLE"
What I'm missing?

Since seems to be a TF bug, I found a workaround to trigger on_epoch_end() of my generator.
class CallbackOnEpochEnd(Callback):
def __init__(self, generator):
super(CallbackOnEpochEnd, self).__init__()
self.generator = generator
def on_epoch_end(self, epoch, logs=None):
self.generator.on_epoch_end()
[...]
generator = HDF5Generator()
d = tf.data.Dataset.from_generator(lambda: generator, output_signature=(tf.TensorSpec(shape=(5,20)), tf.TensorSpec(shape=(1,))))
[...]
on_epoch_end_callback = CallbackOnEpochEnd(generator)
[...]
model.fit(d, epochs=5, callbacks=[on_epoch_end_callback])
With this "RE-SHUFFLE" appear on console after every epoch!

How to implement binary mask matrix in Keras?

I'm currently working on a project and part of it is reimplementing a model written for a paper in PyTorch in Keras. The overal model classifies proteins based on three elements of their properties: sequence, interaction with other proteins, and domains in their sequence (motifs). The part I'm working on recreating currently is the Protein-Protein Interaction part. Firstly, the input vectors simply go through some fully connected layers which is easy enough to implement in keras. However, the outputs from this model are fed into a 'weight classifier model' which applies a binary mask matrix to inputs using a layer created specifically for this model using PyTorch's nn.functional API.
Here is the code I am struggling to implement in keras:
class Weight_classifier(nn.Module):
def __init__(self, func):
super(Weight_classifier, self).__init__()
# self.weight_layer = nn.Linear(OUT_nodes[func]*3, OUT_nodes[func])
self.weight_layer = MaskedLinear(OUT_nodes[func]*3, OUT_nodes[func], 'data/{}_maskmatrix.csv'.format(func)).cuda()
self.outlayer= nn.Linear(OUT_nodes[func], OUT_nodes[func])
def forward(self, weight_features):
weight_out = self.weight_layer(weight_features)
# weight_out = F.sigmoid(weight_out)
weight_out = F.relu(weight_out)
weight_out = F.sigmoid(self.outlayer(weight_out))
return weight_out
class MaskedLinear(nn.Linear):
def __init__(self, in_features, out_features, relation_file, bias=True):
super(MaskedLinear, self).__init__(in_features, out_features, bias)
mask = self.readRelationFromFile(relation_file)
self.register_buffer('mask', mask)
self.iter = 0
def forward(self, input):
masked_weight = self.weight * self.mask
return F.linear(input, masked_weight, self.bias)
def readRelationFromFile(self, relation_file):
mask = []
with open(relation_file, 'r') as f:
for line in f:
l = [int(x) for x in line.strip().split(',')]
for item in l:
assert item == 1 or item == 0 # relation 只能为0或者1
mask.append(l)
return Variable(torch.Tensor(mask))
And this is the paper I am working to, it contains several diagrams and explanations of the models if I have not explained the issue sufficiently.
Many thanks.

How to use multiple models in Keras model subclass API

I am trying to train a fairly complex model that uses multiple frozen pre-trained models and has a custom training loop with a fairly complicated multi-task loss function. Because of these complexities, my plan was to define multiple separate Keras models within the subclassed model. I have been having problems with my setup and I've been able to simplify it to a simple example that demonstrates the problem.
The code below trains a simple model called MainModel, which uses the Keras model subclassing API, but it is basically just a Sequential([Conv1d(), Conv1d()]) model. When I define another model in the same class, self.aux_model, the original model no longer trains properly. In the example, self.aux_model doesn't play any role in the training, it is only defined, never used. Specifically, after each training iteration, the weight values are the same as they were at the beginning of the iteration. So, the model weights are never updated, even though the gradients have non-zero values.
import numpy as np
import tensorflow as tf
from tensorflow.python.keras.callbacks import Callback
num_epochs = 5
steps_per_epoch = 100
audio_len = 16000
class WeightChecker:
"""Automated health checks for training Keras models."""
def __init__(self, model):
self.initial_model = model
self.var_names = [var.name for var in model.trainable_variables]
self.prev_weights = model.get_weights()
def check_epoch(self, model):
"""Checks to run at the end of an epoch"""
self.check_untrained_params(model)
def check_untrained_params(self, model):
"""Compare self.model.trainable_variables to self.prev_weights"""
passed = True
curr_weights = model.get_weights()
for curr_var, prev_var, var_name in zip(curr_weights, self.prev_weights, self.var_names):
eq = np.equal(curr_var, prev_var).all()
if eq:
passed = False
print(f"\nWarning: Variable {var_name} was not updated with training. "
f"Confirm that this layer is correctly "
f"connected to the computation graph.")
self.prev_weights = [w.copy() for w in curr_weights]
return passed
class WeightCheckerCallback(Callback):
"""Check model initialization and run training checks.
"""
def __init__(self):
super().__init__()
self.weight_check = None
def setup_weight_checker(
self,
model: tf.keras.Model = None):
"""Initialize the callback with an input_batch and targets."""
self.weight_check = WeightChecker(model)
def on_train_begin(self, logs=None):
if self.weight_check is None:
raise ValueError("setup_weight_checker() must be called to use WeightCheckerCallback.")
def on_epoch_end(self, epoch, logs=None):
self.weight_check.check_epoch(self.model)
class MainModel(tf.keras.Model):
"""Main Model."""
def __init__(self):
super().__init__()
self.feature_dim = 128
self.aux_model = self._set_aux_model()
self.map_model = tf.keras.Sequential([tf.keras.layers.Conv1D(
64, 3, padding='same'
),
tf.keras.layers.Conv1D(
1, 3, padding='same'
)])
def call(self, inputs, training=True):
output = self.map_model(inputs)
return output
def train_step(self, data):
mixed_audio = data[0]
clean_audio = data[1]
with tf.GradientTape() as tape:
decoded_audio = self.map_model(mixed_audio)
total_loss = tf.reduce_mean(tf.abs(decoded_audio - clean_audio))
grads = tape.gradient(total_loss, self.trainable_variables)
self.optimizer.apply_gradients(zip(grads, self.trainable_variables))
losses = {
'loss': total_loss,
}
return losses
#staticmethod
def _set_aux_model():
"""Set an auxiliary model."""
model = tf.keras.Sequential([tf.keras.layers.Dense(1)])
model.build(input_shape=(None, 1))
model.trainable = False
return model
class TrainingTask:
"""A Keras model training task."""
def __init__(self):
self.model, self.stateful_model = self._set_model()
self.callbacks = [WeightCheckerCallback()]
#staticmethod
def _set_model():
model = MainModel()
# Build the model with fake data.
model.compile(optimizer='adam')
fake_data = np.random.randn(1,
audio_len,
1)
fake_data = fake_data.astype(np.float32)
model(fake_data, training=True)
return model, None
def fit(self):
"""Custom model fit method."""
try:
weight_checker_callback_index = [isinstance(cb, WeightCheckerCallback)
for cb in self.callbacks].index(True)
except ValueError:
weight_checker_callback_index = None
if weight_checker_callback_index is not None:
self.callbacks[weight_checker_callback_index].setup_weight_checker(
model=self.model
)
for callback in self.callbacks:
callback.set_model(self.model)
print("\nBegin training")
for callback in self.callbacks:
callback.on_train_begin()
for epoch in range(num_epochs):
for callback in self.callbacks:
callback.on_epoch_begin(epoch)
for batch in range(steps_per_epoch):
x, y = next(data_gen_batch())
for callback in self.callbacks:
callback.on_batch_begin(batch)
metrics = self.model.train_step([x, y])
batch_loss = np.mean(metrics.pop('loss'))
print(batch, epoch, batch_loss)
for callback in self.callbacks:
callback.on_batch_end(batch, metrics)
print(f'Epoch: {epoch}')
numeric_metrics = dict()
numeric_metrics['loss'] = batch_loss
for callback in self.callbacks:
callback.on_epoch_end(epoch, numeric_metrics)
def data_gen():
"""Generate random data for training."""
data = (np.random.random((audio_len, 1)), np.random.random((audio_len, 1)))
while True:
yield data
def data_gen_batch(batch_size=8):
"""Generate random data in batches for training."""
data = next(data_gen())
data_batch = (np.stack([data[0]] * batch_size, axis=0),
np.stack([data[1]] * batch_size, axis=0))
while True:
yield data_batch
if __name__ == '__main__':
task = TrainingTask()
task.fit()
The classes WeightCheckerCallback and WeightChecker are a callback that I defined to illustrate the problem, which would otherwise result in a silent failure. In addition to some output from each training step, the code will produce the following warnings about layers of the map_model, which should be updating (aux_model only has a Dense layer):
Warning: Variable main_model/sequential_1/conv1d/kernel:0 was not updated with training. Confirm that this layer is correctly connected to the computation graph.
Warning: Variable main_model/sequential_1/conv1d/bias:0 was not updated with training. Confirm that this layer is correctly connected to the computation graph.
However, if the aux_model is commented out, the warnings will not appear and the model weights will be updated as expected.
# self.aux_model = self._set_aux_model()
Obviously, there are several ways in tensorflow to get this simple Sequential model to train properly, so I'm not just looking for a workaround to get this particular example working. Rather, I'm hoping that someone can explain what is going on with this example in terms of the Tensorflow sessions and graphs involved, as well as what the best practices are for avoiding conflicts between multiple different Keras models when nesting them with the subclass API. My ultimate goal is to train a more complex system of models using a similar framework.

This turned out to be a problem with the way that I was checking the layers with the WeightChecker class. model.get_weights() returns all the weights, not just the trainable weights. Therefore, when we use zip() in the for loop, we are zipping together lists with different lengths, and this causes the name of the layer that's not being updated to be misreported. The bug can be solved by using the following rather than model.get_weights():
self.prev_weights = [var.numpy() for var in model.trainable_variables]

Creating a new object every time you call a function in python

I have a recommender system that I need to train, and I included the entire training procedure inside a function:
def train_model(data):
model = Recommender()
Recommender.train(data)
pred = Recommender.predict(data)
return pred
something like this. Now if I want to train this inside a loop, for different datasets, like:
preds_list = []
data_list = [dataset1, dataset2, dataset3...]
for data_subset in data_list:
preds = train_model(data_subset)
preds_list += [preds]
How can I make sure that every time I call the train_model function, a brand new instance of a recommender is created, not an old one, trained on the previous dataset?

You are already creating a new instance everytime you execute train_model. The thing you are not using the new instance.
You probably meant:
def train_model(data):
model = Recommender()
model.train(data)
pred = model.predict(data)
return pred

Use the instance you've instantiated, not the class
class Recommender:
def __init__(self):
self.id = self
def train(self, data):
return data
def predict(self, data):
return data + str(self.id)
def train_model(data):
model = Recommender()
model.train(data)
return model.predict(data)
data = 'a data '
x = {}
for i in range(3):
x[i] = train_model(data)
print(x[i])
# a data <__main__.Recommender object at 0x11cefcd10>
# a data <__main__.Recommender object at 0x11e0471d0>
# a data <__main__.Recommender object at 0x11a064d50>