Adam with weight decay per parameter + some activation functions #452
Conversation
|
|
||
| class SiLU(Module): | ||
| r""" | ||
| SiLU activation function. Also known as the "swish" function. |
There was a problem hiding this comment.
Could you document the expression here as well?
drjit/opt.py
Outdated
| beta_2: float = 0.999, | ||
| epsilon: float = 1e-8, | ||
| weight_decay: float = 0.01, | ||
| weight_decay: Optional[float | Mapping[str, float]] = None, |
There was a problem hiding this comment.
weight decay is an AdamW-specific feature. I would prefer to keep this 100% in the AdamW subclass so that we don't have to pay any cost (like unpacking above) in the general optimizer.
When optimizing simple calculations (e.g. neural nets) without any ray tracing, the Python code in the optimizer can actually inhibit saturating the GPU, hence this focus on keeping the code here as simple/efficient as possible. It was also the motivation for a larger rewrite when moving the optimizer code from Mitsuba to Dr.Jit.
AdamW can store an optional per-parameter weight decay override using the extra field (setting this value to None by default).
There was a problem hiding this comment.
I did not get your comment here. The other optimizers do not have any access to the weight_decay parameter. I just have an additional dictionary in AdamW and increase the size of the tuple (extra) inset_weight_decay function. Can you elaborate your comment about keeping everything in AdamW?
There was a problem hiding this comment.
Sorry, I misread the diff and had thought that you were modifying the base optimizer. I think it would be better if the set_weight_decay function works like the set_learning_rate function, please check out the docstring of this function. Right now, all of the optimizers are written so that they don't use key-based dictionary lookups in opt.step(). They iterate over dictionaries but never explicitly search a dictionary for a string-based key (which involves string hashing etc.)
|
FYI: The |
Some activation functions which are helpful for neural network training such as Sigmoid, Softplus and Swish. The AdamW is extended to handle per parameter weight decays.