Recording Model Topology

Hi everyone, relatively new to WandB here. From my understanding, I can log the model topology so like each of the layers, how many neurons, kernel, etc. by using the .watch() command. So for example I am making a basic GAN with a discriminator and generator using nn.Sequential like this:

class Discriminator(nn.Module):
def init(self):
super(Discriminator, self).init()
self.model = nn.Sequential(
nn.Linear(784, 512),
nn.BatchNorm1d(512),
nn.LeakyReLU(0.2),
nn.Dropout(),
nn.Linear(512, 256),
nn.LeakyReLU(0.2),
nn.Dropout(),
nn.Linear(256, 1),
nn.Sigmoid()

Then I have my GAN class

class GAN(pl.LightningModule):
def init(self,noise_dim=100,lr=0.0002):
super().init()

where I have these lines to define the generator and discriminator and log them

    self.generator=Generator(noise_dim=self.hparams.noise_dim)
    self.discriminator=Discriminator()
    
    self.generator.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
    self.discriminator.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))

    wandb.watch(self.generator, log='all', log_freq=100)
    wandb.watch(self.discriminator, log='all', log_freq=100)

Now when I open the WandB page to check my runs, i can see that it is logging the gradients and weights and biases but I cannot find the model topology describing the layers I used. Is there something I am doing wrong?

Thank you in advance!

Hi @8688chris , thanks for reaching out with your question. Would you mind sharing the URL for the Run in question and the rest of the code you are running to log the data to W&B?

Hi there! Thanks for the reply. Here is the workspace URL Weights & Biases

There are many runs but I do not change the logging procedure. The main thing is not being able to see the model topology, maybe one of my logging lines is missing something.

and here is a link to my .ipynb which I put in colab

Thank you for sharing the URL and the code - I will have a look at these and will get back to you with additional information, most likely next week.

Hi @8688chris , thank you for your patience here.

Having a look at this, the .watch method would need to be called on the instantiated WandbLogger for the topology to be stored on W&B.

See this example:

# setup model
model = LitMLP(in_dims=(1, 28, 28))

wandb_logger = WandbLogger(entity=entity,
                           project=project,
                           log_model='all')

wandb_logger.watch(model,log='all')

trainer = pl.Trainer(
    logger=wandb_logger,    # W&B integration
    log_every_n_steps=50,   # set the logging frequency
    max_epochs=5,           # number of epochs
    )
# fit the model
trainer.fit(model, mnist)

wandb.finish()

Which generated this Run

Ohhh ok, I did not know i needed to .watch the logger itself. Thank you so much for your reply
I really appreciate it!

You are welcome! I will mark this request as solved now - please feel free to reach out in the future for any further questions.