search

lucidrains / imagen-pytorch

Implementation of Imagen, Google's Text-to-Image Neural Network, in Pytorch
MIT License
8.09k stars 768 forks source link

Visualization of the network #349

Open kirilllzaitsev opened 1 year ago

kirilllzaitsev commented 1 year ago

Hi, do you have in mind a tool to visualize a sample Imagen network from the main README?

torchsummary, onnx don't work out-of-the-box, since forward passes within Imagen are quite complex in terms of their args and kwargs.