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]
Related
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.
The topic above is a bit ambiguous, the explaination is below:
class Trainer:
"""Object used to facilitate training."""
def __init__(
self,
# params: Namespace,
params,
model,
device=torch.device("cpu"),
optimizer=None,
scheduler=None,
wandb_run=None,
early_stopping: callbacks.EarlyStopping = None,
):
# Set params
self.params = params
self.model = model
self.device = device
# self.optimizer = optimizer
self.optimizer = self.get_optimizer()
self.scheduler = scheduler
self.wandb_run = wandb_run
self.early_stopping = early_stopping
# list to contain various train metrics
# TODO: how to add more metrics? wandb log too. Maybe save to model artifacts?
self.history = DefaultDict(list)
#staticmethod
def get_optimizer(
model: models.CustomNeuralNet,
optimizer_params: global_params.OptimizerParams(),
):
"""Get the optimizer for the model.
Args:
model (models.CustomNeuralNet): [description]
optimizer_params (global_params.OptimizerParams): [description]
Returns:
[type]: [description]
"""
return getattr(torch.optim, optimizer_params.optimizer_name)(
model.parameters(), **optimizer_params.optimizer_params
)
Notice that initially I passed in optimizer in the constructor, where I will be calling it outside this class. However, I now put get_optimizer inside the class itself (for consistency purpose, but unsure if it is ok). So, should I still define self.optimizer = self.get_optimizer() or just use self.get_optimizer() at the designated places in the class? The former encourages some readability for me.
Addendum: I now put the instance inside the .fit() method where I will call say 5 times to train the model 5 times. In this scenario, even though there won't be any obvious issue as we are using optimizer once per call, will it still be better to not define self.optimizer here?
def fit(
self,
train_loader: torch.utils.data.DataLoader,
valid_loader: torch.utils.data.DataLoader,
fold: int = None,
):
"""[summary]
Args:
train_loader (torch.utils.data.DataLoader): [description]
val_loader (torch.utils.data.DataLoader): [description]
fold (int, optional): [description]. Defaults to None.
Returns:
[type]: [description]
"""
self.optimizer = self.get_optimizer(
model=self.model, optimizer_params=OPTIMIZER_PARAMS
)
self.scheduler = self.get_scheduler(
optimizer=self.optimizer, scheduler_params=SCHEDULER_PARAMS
)
There is a difference between the two: calling your get_optimizer will instantiate a new torch.optim.<optimizer> every time. In contrast, setting self.optimizer and accessing it numerous times later will only create a single optimizer instance.
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 can I get the value of the learning rate updated at each on_train_batch_begin?
lr_decayed_fn = tf.keras.experimental.CosineDecay(initial_lr, decay_steps)
model.compile(optimizer=tf.keras.optimizers.SGD(learning_rate=lr_decayed_fn))
I tried this way
def on_train_batch_begin (self, batch, logs = None):
lr = K.get_value(self.model.optimizer.lr)
but I get <tensorflow.python.keras.optimizer_v2.learning_rate_schedule.CosineDecay object at 0x7f ...>
When you set a function as a learning rate or an object subclassing LearningRateScheduler, you need to call that function (or Callable) with the current training step to get the learning rate. You can get the current training step by using the iterations attribute of the optimizer.
class CustomCallback(tf.keras.callbacks.Callback):
def __init__(self) -> None:
super().__init__()
def on_train_batch_begin(self, batch, logs=None):
lr = tf.keras.backend.get_value(
self.model.optimizer.lr(self.model.optimizer.iterations)
)
I want to use model.predict in batch generator, what a possible ways of achieve this?
Seems one option is to load model on init and on epoch end:
class DataGenerator(keras.utils.Sequence):
def __init__(self, model_name):
# Load model
# ...
def on_epoch_end(self):
# Load model
In my experience, predicting another model while training will bring errors.
You should probably simply append your training model after your generator model.
Suppose you have:
generator_model (the one you want to use inside the generator)
training_model (the one you want to train)
Then
generatorInput = Input(shapeOfTheGeneratorInput)
generatorOutput = generator_model(generatorInput)
trainingOutput = training_model(generatorOutput)
entireModel = Model(generatorInput,trainingOutput)
Make sure that the generator model has all layers untrainable before compiling:
genModel = entireModel.layers[1]
for l in genModel.layers:
l.trainable = False
entireModel.compile(optimizer=optimizer,loss=loss)
Now use the generator regularly.
Predicting inside the generator:
class DataGenerator(keras.utils.Sequence):
def __init__(self, model_name, modelInputs, batchSize):
self.genModel = load_model(model_name)
self.inputs = modelInputs
self.batchSize = batchSize
def __len__(self):
l,rem = divmod(len(self.inputs), self.batchSize)
return (l + (1 if rem > 0 else 0))
def __getitem__(self,i):
items = self.inputs[i*self.batchSize:(i+1)*self.batchSize]
items = doThingsWithItems(items)
predItems = self.genModel.predict_on_batch(items)
#the following is the only reason not to chain models
predItems = doMoreThingsWithItems(predItems)
#do something to get Y_train_items as well
return predItems, y_train_items
If you do find the error I mentioned, you can sacrifice the parallel generation capabilities and do some manual loops:
for e in range(epochs):
for i in range(batches):
x,y = generator[i]
model.train_on_batch(x,y)