IvanDrokin / torch-conv-kan

This project is dedicated to the implementation and research of Kolmogorov-Arnold convolutional networks. The repository includes implementations of 1D, 2D, and 3D convolutions with different kernels, ResNet-like and DenseNet-like models, training code based on accelerate/PyTorch, as well as scripts for experiments with CIFAR-10 and Tiny ImageNet.
MIT License
416 stars 32 forks source link

output_hook is not storing tensors #1

Closed JiwonKKim closed 6 months ago

JiwonKKim commented 6 months ago

Hi. thanks for your awesome source code.

I'm training the model with your example code with l1 & l2 regularization. I tried to check the gradient flow of the tensor during the backward process. However, I noticed that in train/trainer.py, output_hook was not storing the gradient tensors. Seems like the registering hook function is not working in train/trainer.py, line 186.

Should I fix the code from

for module in model.named_modules():

to

for name, module in model.named_modules():

in train/trainer.py, line 184? Or did I miss something?

My torch version is 2.1

Thanks.

IvanDrokin commented 6 months ago

Hi there! Thanks for your feedback, I'll take a look at this and return soon.

IvanDrokin commented 6 months ago

@JiwonKKim Hi again!

One more time - many thanks for your feedback, much appreciate it

You are right, the output_hook doesn't work, because there are no named modules.

The correct way to fix it is to change

for module in model.named_modules():

to

for module in model.modules():

I'll update my repository today later, Have a good day!

IvanDrokin commented 6 months ago

@JiwonKKim I've committed a fix for the issue, please check out my latest commit. Please, feel free to ask any question you have.

JiwonKKim commented 6 months ago

Thanks!