Closed AliOskooeiTR closed 3 years ago
@AliOskooeiTR Hi Ali, it is tricky because in the encoder / decoder setting, the hook for updating the kmeans must be moved from the individual transformers to the EncDec class, as shown https://github.com/lucidrains/routing-transformer/blob/master/routing_transformer/encoder_decoder.py#L78 Let me work on this later this weekend and see if there is a clean way to do this
@AliOskooeiTR also, as shown https://github.com/lucidrains/routing-transformer/blob/master/routing_transformer/encoder_decoder.py#L59
@AliOskooeiTR if you want to try version 1.4.3, this should fix it
enc.cancel_kmeans_update() # cancel kmeans update for enc and dec
dec.cancel_kmeans_update()
encdec.enc = enc
encdec.dec = dec
encdec.register_kmeans_update() # reregister kmeans
@lucidrains I tried this:
pretrained_lm.cancel_kmeans_update() # cancel kmeans update for enc and dec
enc_dec.enc = pretrained_lm
enc_dec.dec = AutoregressiveWrapper(pretrained_lm)
enc_dec.register_kmeans_update() # reregister kmeans
And unfortunately I still get the same error:
RuntimeError: new kmeans has not been supplied
@AliOskooeiTR ohh, I see the problem, I don't think you can use the same pretrained LM for both encoder and decoder. Instead of doing it that way, why not do something like
enc_dec.enc.load_state_dict(pretrained_lm.state_dict())
enc_dec.dec.load_state_dict(pretrained_lm.state_dict())
@lucidrains Thanks for the suggestion. I tried this:
enc_dec.enc.load_state_dict(pretrained_lm.state_dict())
enc_dec.dec.load_state_dict(pretrained_lm.state_dict())
It would give me missing key error as the state dict keys for enc_dec.enc and pretrained_lm were different. Surprisingly even the number of keys were quite different. I did not understand how this could be since enc_dec.enc and pretrained_lm are both RoutingTransformerLMs.
I then replaced encdec state dicts just to see what happens:
pretrained_lm.cancel_kmeans_update()
pretrained_dec = AutoregressiveWrapper(pretrained_lm)
enc_dec.enc.state_dict=pretrained_lm.state_dict
enc_dec.dec.state_dict=pretrained_dec.state_dict
enc_dec.register_kmeans_update()
Trying to train this gave me the same kmeans error:
RuntimeError: new kmeans has not been supplied
Just closing this issue as I have figured out why it wasn't possible to use the same LM for both the encoder and the decoder. The decoder LM must receive context and be causal. This results in the encoder having a different architecture and state dictionary and not exchangeable with the encoder LM.
I am trying to build and train an encoder-decoder from pretrained routing transformer LMs. The way I approached it was to replace the encoder and decoder in a RoutingTransformerEncDec with the pre-trained RoutingTransformerLMs as follows:
and then try to train the enc_dec as normal when I get the following error:
I would appreciate any feedback on what may be the problem or what is the best way to build an enc_dec from pretrained LM checkpoints.