question about dragonnet

[guanying-zhu]

Hello, quick question about dragonnet. Line 209 of dragonnet.py, if Adam is selected, it seems that the model would first fit using Adam and then fit with SGD. Is that intended?

[0trm]

Yes, that's intended. It follows the Dragonnet authors' own training code: in ihdp_main.py the model is compiled with Adam (lr 1e-3, early stopping after 2 epochs without improvement), fit, then recompiled with SGD with Nesterov momentum (lr 1e-5, patience 40) and fit again. causalml uses the same two phases and callbacks (dragonnet.py), only with fewer default epochs (30 Adam and 100 SGD instead of 100 and 300).

The Adam phase is optional, and each phase has its own arguments:

from causalml.inference.tf import DragonNet

dragon = DragonNet(use_adam=False)   # SGD only
dragon = DragonNet(adam_epochs=30, adam_learning_rate=1e-3,   # Adam phase
                   epochs=100, learning_rate=1e-5, momentum=0.9)   # SGD phase

The JAX version added in 0.17.0 (from causalml.inference.jax import DragonNet) trains the same way and takes the same arguments; with verbose=True it logs [Adam] epochs followed by [SGD] epochs, and only [SGD] with use_adam=False. Checked on causalml 0.17.0.