Integration of the WeightCheckpoint Metric

[OliviaGraeupner]

Hello everybody,

I'm using the following basic code for implementing an continual learning model:

benchmark = SplitMNIST(n_experiences=5, return_task_id=True, seed=42, shuffle=True)
train_stream = benchmark.train_stream
test_stream = benchmark.test_stream
model = SimpleMLP(num_classes=benchmark.n_classes)
optimizer = SGD(model.parameters(), lr=0.001, momentum=0.9)
criterion = CrossEntropyLoss()

loggers = [] 
loggers.append(TextLogger(open('log.txt', 'a')))
loggers.append(InteractiveLogger())
loggers.append(TensorboardLogger()) 
loggers.append(WandBLogger(project_name="avalanche", run_name="run1"))

early_stopping = EarlyStoppingPlugin(patience=10, val_stream_name='train')
replay = ReplayPlugin(mem_size=100)
ewc = EWCPlugin(ewc_lambda=0.001)

eval_plugin = EvaluationPlugin(
    WeightCheckpoint(),
    loggers=loggers,
    collect_all=True,
    benchmark=benchmark,
    strict_checks=False)

cl_strategy = Naive(
    model, optimizer, criterion,
    plugins=[early_stopping, replay, ewc],
    train_mb_size=100, train_epochs=4, eval_mb_size=100, eval_every=1,
    evaluator=eval_plugin)

 print("Starting experiment...")
    results = []
    for experience in benchmark.train_stream:
        print("Start of experience: ", experience.current_experience)
        print("Current Classes: ", experience.classes_in_this_experience)
        cl_strategy.train(experience, eval_streams=[benchmark.test_stream])
        print("Training completed")
        results.append(cl_strategy.eval(benchmark.test_stream))

My goal is to log the weights at certain checkpoints during the training. (The performance of the model is incidental.) For this purpose I plan to use the .update() and .result() functions so I can log the weights of all layers of the SimpleMLP model. My problem is that I don't quite know where to insert the WeightCheckpoint functions in my code appropriately and how to output the results.

weight_checkpoint = WeightCheckpoint()
weight_checkpoint_init = weight_checkpoint.result() # output: None
weight_checkpoint.update(model.features[0].weight.data) # How to update weight_checkpoint automatically for each layer?
weight_checkpoint = weight_checkpoint.result()
...

Accordingly there are no values in the metric_dict.

metric_dict = eval_plugin.get_all_metrics()

In addition, I would like to save the weights in the log.txt file using the TextLogger. How can I integrate this feature in the TextLogger source code?

Thanks for your help!

[AntonioCarta]

The metrics and logging system works best with scalar value. To save a weight checkpoint, I suggest to use a normal plugin, something like:

class WeightCheckpoint:
    def after_training_epoch(...):
        # put the condition to save the model and store in the file system here.

Depending on when and how often you need to save the model, you can use a different callback (e.g. after_training_iteration or after_training_exp).

[OliviaGraeupner]

Thanks for your reply. I don't really understand your answer. Maybe I'll rephrase my question. I know that there is a predefined metric called WeightCheckpoint. I don't want to edit this metric definition. My question was, how can I include this metric in my programme code so that the weights are saved at predefined checkpoints. Furthermore I want to output the saved weights. How can I realize this?