Can't pickle local object with Pipe and ConvNet about revert

[opeltre]

Quelques tests:

1. Fonctionne:

>>> from revert.models import BarlowTwins
>>> from revert.models import ConvNet
>>> model = ConvNet([[10, 4, 2],[12, 8, 3]])
>>> twins = BarlowTwins(model)
>>> twins.save("twins.pt")
>>> BarlowTwins.load("twins.pt")
BarlowTwins(
  (model): ConvNet(
    (dropout): Dropout(p=0.0, inplace=False)
    (conv0): Conv1d(4, 8, kernel_size=(2,), stride=(1,), padding=same, padding_mode=circular)
    (pool0): MaxPool1d(kernel_size=0, stride=0, padding=0, dilation=1, ceil_mode=False)
  )
)

Ici la classe BarlowTwins définit sa loss comme méthode

2. Fonctionne:

>>> from revert.models import ConvNet, BarlowTwins
>>> f = ConvNet([[10, 5, 5], [10, 5, 5]])
>>> g = ConvNet([[10, 5, 5], [10, 5, 5]])
>>> h = f @ g
>>> h.save("pipe.pt")
>>> from revert.models import Module
>>> Module.lo
Module.load(             Module.load_state_dict(  Module.loss_on(
>>> Module.load("pipe.pt")
Pipe(
  (module0): ConvNet(
    (dropout): Dropout(p=0.0, inplace=False)
    (conv0): Conv1d(5, 5, kernel_size=(5,), stride=(1,), padding=same, padding_mode=circular)
    (pool0): MaxPool1d(kernel_size=1, stride=1, padding=0, dilation=1, ceil_mode=False)
  )
  (module1): ConvNet(
    (dropout): Dropout(p=0.0, inplace=False)
    (conv0): Conv1d(5, 5, kernel_size=(5,), stride=(1,), padding=same, padding_mode=circular)
    (pool0): MaxPool1d(kernel_size=1, stride=1, padding=0, dilation=1, ceil_mode=False)
  )
)

Ici pas de loss définie

3. Ne fonctionne pas:

>>> h.loss = lambda y, yt: (y - yt).norm()
>>> h.save("pipe.pt")
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/oli/revert/revert/models/module.py", line 86, in save
    torch.save(self, path)
  File "/home/oli/.local/lib/python3.9/site-packages/torch/serialization.py", line 379, in save
    _save(obj, opened_zipfile, pickle_module, pickle_protocol)
  File "/home/oli/.local/lib/python3.9/site-packages/torch/serialization.py", line 484, in _save
    pickler.dump(obj)
_pickle.PicklingError: Can't pickle <function <lambda> at 0x7f4922901280>: attribute lookup <lambda> on __main__ failed

Le problème semble venir du fait que python ne sait pas sérialiser une fonction et que la classe en elle-même ne définit pas loss

[opeltre]

Du coup il faudrait interdire les cas où l'on ajoute la loss "à la main", même avec un setattr(...)

Il faudrait définir une classe à chaque fois que l'on veut changer de loss, par exemple:

class ConvNet_MSE(ConvNet): 

    def loss(self, y, y_tgt): 
        return ((y - y_tgt)**2).mean()

[EDIT] pas obligé

[opeltre]

Il faut aussi redéfinir la méthode loss de la classe Pipe comme ça: (et supprimer la partie correspondante du __init__)

class Pipe(Module):

    def loss(self, y, *ys): 
        return self.modules[-1].loss(y, *ys)

[opeltre]

N.B. pour pas trop s'embêter on pourra aussi définir une loss par défaut dans Module

[opeltre]

Erratum:

Celui-ci fonctionne (4):

>>> def mseLoss(y, yt) : return ((y - yt)**2).mean()
>>> h = f @ g
>>> setattr(h, "loss", mseLoss)
>>> h.save("pipe.pt")
>>> h = Module.load("pipe.pt")
>>> h.loss
<function mseLoss at 0x7f....>

Donc vraisemblablement, si la fonction est définie dans le scope local, pickle va seulement sauvegarder le nom de la fonction pour tenter de de la retrouver dans le scope en lecture

=> pas possible pour les lambdas (anonymes)

Ca marche bien si on ferme la session et retente de charger h, pourvu qu'une fonction mseLoss existe dans le scope 🙂

Du coup il suffirait de garder nos loss définies dans un module à part!

[opeltre]

Autre problème était: sérialisation de module.writer avant qu'il ne soit fermé

Tout marche bien après appel de module.writer.close()