Torch Parallel Layers Visualization in wandb

username: ColdFrenzy
cooked: <p>Hi everyone,<br>
I created a dynamic actor-critic module deriving from nn.Module. The module is made up of 3 submodules: a shared network made up of a number of convolutional layers and 2 independent parts made up of fc layers that receive the flattened output of the previous module as input. The code that I used is the following:</p>
<pre><code class="lang-auto">class Conv3DModelFree(nn.Module):
    def __init__(self, in_shape, num_actions, **kwargs):
        fc_layers = kwargs["fc_layers"]
        super(Conv3DModelFree, self).__init__(
            num_actions, features_out=fc_layers[-1])
        # super().__init__()
        conv_layers = kwargs["conv_layers"]

        self.shared_layers = []
        self.num_actions = num_actions
        self.in_shape = in_shape
        self.num_channels = in_shape[0]
        self.num_frames = kwargs["num_frames"]
        next_inp = None
        # =============================================================================
        # FEATURE EXTRACTOR SUBMODULE
        # =============================================================================
        for i, cnn in enumerate(conv_layers):
            if i == 0:
                self.shared_layers.append(nn.Conv3d(
                    self.num_channels, cnn[0], kernel_size=cnn[1], stride=cnn[2]))
                self.shared_layers.append(nn.LeakyReLU())
            else:
                self.shared_layers.append(nn.Conv3d(
                    next_inp, cnn[0], kernel_size=cnn[1], stride=cnn[2]))
                self.shared_layers.append(nn.LeakyReLU())
            next_inp = cnn[0]

        for i, layer in enumerate(self.shared_layers):
            if i == 0:
                fake_inp = torch.zeros(
                    [1, self.num_channels, self.num_frames, *self.in_shape[1:]])
                fake_inp = self.shared_layers[i](fake_inp)
            else:
                fake_inp = self.shared_layers[i](fake_inp)
        next_inp = fake_inp.view(1, -1).size(1)

        # =============================================================================
        # ACTOR AND CRITIC SUBMOODULES
        # =============================================================================
        self.actor_subnet = []
        self.critic_subnet = []
        for i, fc in enumerate(fc_layers):
            if i == 0:
                self.shared_layers.append(nn.Linear(next_inp, fc))
                self.shared_layers.append(nn.LeakyReLU())
            else:
                # Separate submodules for the actor and the critic
                self.actor_subnet.append(nn.Linear(next_inp, fc))
                self.critic_subnet.append(nn.Linear(next_inp, fc))
                self.actor_subnet.append(nn.LeakyReLU())
                self.critic_subnet.append(nn.LeakyReLU())
            next_inp = fc
def forward(self, input)
        for i, layer in enumerate(self.shared_layers[:-2]):
            if i == 0:
                x = self.shared_layers[i](input)
            else:
                x = self.shared_layers[i](x)
        # last 2 shared layers requires a reshape of the input
        x = x.view(x.shape[0], -1)
        x = self.shared_layers[-2](x)
        x = self.shared_layers[-1](x)

        action_logits = None
        for i, layer in enumerate(self.actor_subnet):
            if i == 0:
                action_logits = self.actor_subnet[i](x)
            else:
                action_logits = self.actor_subnet[i](action_logits)
        action_logits = self.actor(action_logits)

        value = None
        for i, layer in enumerate(self.critic_subnet):
            if i == 0:
                value = self.critic_subnet[i](x)
            else:
                value = self.critic_subnet[i](value)
        value = self.critic(value)

        return action_logits, value
</code></pre>
<p>When I use wandb.watch on this class, the graphs do not appear to me probably because wandb cannot identifies the submodules of the network that are included in a list. I tested a different solution, trying to wrap the various pieces in some nn.Sequential modules as follows:</p>
<pre><code class="lang-auto">class Conv3DModelFree(OnPolicy):
    def __init__(self, in_shape, num_actions, **kwargs):
        fc_layers = kwargs["fc_layers"]
        super(Conv3DModelFree, self).__init__(
            num_actions, features_out=fc_layers[-1])
        # super().__init__()
        conv_layers = kwargs["conv_layers"]

        shared_layers = OrderedDict()
        self.num_actions = num_actions
        self.in_shape = in_shape
        self.num_channels = in_shape[0]
        self.num_frames = kwargs["num_frames"]
        next_inp = None
        # =============================================================================
        # FEATURE EXTRACTOR SUBMODULE
        # =============================================================================
        for i, cnn in enumerate(conv_layers):
            if i == 0:
                shared_layers["conv_0"] = nn.Conv3d(
                    self.num_channels, cnn[0], kernel_size=cnn[1], stride=cnn[2])
                shared_layers["activ_0"] = nn.LeakyReLU()
            else:
                shared_layers["conv_"+str(i)] = nn.Conv3d(
                    next_inp, cnn[0], kernel_size=cnn[1], stride=cnn[2])
                shared_layers["activ_" + str(i)] = nn.LeakyReLU()
            next_inp = cnn[0]

        # flatten the output starting from dim=1 by default
        shared_layers["flatten"] = nn.Flatten()

        for i, layer in enumerate(shared_layers):
            if i == 0:
                fake_inp = torch.zeros(
                    [1, self.num_channels, self.num_frames, *self.in_shape[1:]])
                fake_inp = shared_layers[layer](fake_inp)
            else:
                fake_inp = shared_layers[layer](fake_inp)
        next_inp = fake_inp.view(1, -1).size(1)

        # =============================================================================
        # ACTOR AND CRITIC SUBMODULES
        # =============================================================================
        actor_subnet = OrderedDict()
        critic_subnet = OrderedDict()
        for i, fc in enumerate(fc_layers):
            if i == 0:
                shared_layers["fc_0"] = nn.Linear(next_inp, fc)
                shared_layers["fc_activ"] = nn.LeakyReLU()
            else:
                # Separate submodules for the actor and the critic
                actor_subnet["actor_fc_"+str(i)] = nn.Linear(next_inp, fc)
                critic_subnet["critic_fc_" +
                              str(i)] = nn.Linear(next_inp, fc)
                actor_subnet["actor_activ_"+str(i)] = nn.LeakyReLU()
                critic_subnet["critic_activ_"+str(i)] = nn.LeakyReLU()
            next_inp = fc
        actor_subnet["actor_out"] = nn.Linear(next_inp, self.num_actions)
        critic_subnet["critic_out"] = nn.Linear(next_inp, 1)

        self.shared_network = nn.Sequential(shared_layers)
        self.actor = nn.Sequential(actor_subnet)
        self.critic = nn.Sequential(critic_subnet)
    
    def forward(self, input):
     
        shared_net_output = self.shared_network(input)
        action_logits = self.actor(shared_net_output)
        value = self.critic(shared_net_output)

        return action_logits, value
</code></pre>
<p>Using this different implementation, I can visualize the graphs but I get completely different results and extremely high loss values. Does anyone know a correct way to define this architecture so that it is viewable on wandb?</p>

username: thecho7
cooked: <p><strong>First code</strong></p>
<p>The graph cannot be visualized because the layers are not in <em>nn</em> container.<br>
You should put them into a container like below,</p>
<pre><code class="lang-auto">self.nn_shared_layers = nn.ModuleList(*self.shared_layers)
self.nn_actor_subnet = nn.ModuleList(*self.actor_subnet)
self.nn_critic_subnet = nn.ModuleList(*self.critic_subnet)
</code></pre>
<p>I recommend you to use <em>nn.Sequential</em> instead of <em>nn.ModuleList</em>.<br>
Except the <em>shared_layers</em>, your forward function is designed in a serial manner. In short,<br>
You can rewrite your forward function like</p>
<pre><code class="lang-auto"># In __init__
self.nn_actor_subnet = nn.Sequential(*self.actor_subnet)
# In forward
action_logits = self.nn_actor_subnet(x)
</code></pre>
<p><strong>Second Code</strong><br>
I guess your code is well modified to visualize the graph but the reason that your result seems weird comes from not using activation functions.<br>
I don’t know how small your learning rate and which optimizer you use. But if the learning rate is not that big such as 1, the problem is usually mis-designed network.</p>
<hr>
<pre><code class="lang-auto">class Conv3DModelFree(nn.Module):
    def __init__(self, in_shape, num_actions, **kwargs):
        fc_layers = kwargs["fc_layers"]
        super(Conv3DModelFree, self).__init__(
            num_actions, features_out=fc_layers[-1])
        # super().__init__()
        conv_layers = kwargs["conv_layers"]

        self.shared_layers = []
        self.num_actions = num_actions
        self.in_shape = in_shape
        self.num_channels = in_shape[0]
        self.num_frames = kwargs["num_frames"]
        next_inp = None
        # =============================================================================
        # FEATURE EXTRACTOR SUBMODULE
        # =============================================================================
        for i, cnn in enumerate(conv_layers):
            if i == 0:
                self.shared_layers.append(nn.Conv3d(
                    self.num_channels, cnn[0], kernel_size=cnn[1], stride=cnn[2]))
                self.shared_layers.append(nn.LeakyReLU())
            else:
                self.shared_layers.append(nn.Conv3d(
                    next_inp, cnn[0], kernel_size=cnn[1], stride=cnn[2]))
                self.shared_layers.append(nn.LeakyReLU())

            if i == len(conv_layers) - 3:
                self.shared_layers.append(nn.Flatten(start_dim=1))

            next_inp = cnn[0]

        for i, layer in enumerate(self.shared_layers):
            if i == 0:
                fake_inp = torch.zeros(
                    [1, self.num_channels, self.num_frames, *self.in_shape[1:]])
                fake_inp = self.shared_layers[i](fake_inp)
            else:
                fake_inp = self.shared_layers[i](fake_inp)
        next_inp = fake_inp.view(1, -1).size(1)

        # =============================================================================
        # ACTOR AND CRITIC SUBMOODULES
        # =============================================================================
        self.actor_subnet = []
        self.critic_subnet = []
        for i, fc in enumerate(fc_layers):
            if i == 0:
                self.shared_layers.append(nn.Linear(next_inp, fc))
                self.shared_layers.append(nn.LeakyReLU())
            else:
                # Separate submodules for the actor and the critic
                self.actor_subnet.append(nn.Linear(next_inp, fc))
                self.critic_subnet.append(nn.Linear(next_inp, fc))
                self.actor_subnet.append(nn.LeakyReLU())
                self.critic_subnet.append(nn.LeakyReLU())
            next_inp = fc

        self.nn_shared_layers = nn.Sequential(*self.shared_layers)
        self.nn_actor_subnet = nn.Sequential(*self.actor_subnet)
        self.nn_critic_subnet = nn.Sequential(*self.critic_subnet)

    def forward(self, input)
        x = self.nn_shared_layers(input)
        action_logits = self.nn_actor_subnet(x)
        value = self.nn_critic_subnet(x)

        return action_logits, value</code></pre>

username: ColdFrenzy
cooked: <p>Thanks a lot for the answer.<br>
I tried your solution but the network behaves exactly like in my second code.</p>
<p>Actually, the network in the second code contains the activation functions, I checked that the architecture was correct also through tensorboard.<br>
The learning rate is 3e-4 but in all cases, the program is identical for both models I have defined, so there are no other elements outside the network that influence the training.<br>
I don’t know, maybe using 2 Sequential blocks in parallel like in the second code can cause problems with backpropagation (I’m just assuming)</p>

username: anantguptadbl
cooked: <p><a class="mention" href="/u/coldfrenzy">@ColdFrenzy</a></p>
<p>nn.Sequential should work as mentioned by <a class="mention" href="/u/thecho7">@thecho7</a>. Below is an example</p>
<p><strong>ORIGINAL</strong></p>
<pre><code class="lang-auto">import numpy as np
import time
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable

class random_model(nn.Module):
    def __init__(self, num_layers):
        super(random_model, self).__init__()
        self.model1 = nn.Sequential(nn.Linear(100, 20), nn.BatchNorm1d(20), nn.ReLU())
        self.model2 = []
        for i in range(num_layers):
            self.model2.append(nn.Sequential(nn.Linear(20, 20), nn.BatchNorm1d(20), nn.ReLU()))
        #self.model2 = nn.Sequential(*self.model2)
        self.model3 = nn.Linear(20, 1)
    
    def forward(self, x):
        x = self.model1(x)
        for cur_model in self.model2:
            x = cur_model(x)
        #x = self.model2(x)
        x = self.model3(x)
        return x
    
model = random_model(10)
loss = torch.nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

X = torch.rand(100, 100)
y = torch.rand(100)

import wandb
wandb.init(project="random-project-custom-class", entity="")
wandb.config = {
  "learning_rate": 0.001,
  "epochs": 100,
  "batch_size": 100
}

# 3. Log gradients and model parameters
wandb.watch(model, log="gradients", log_freq=100)
print("DEF")
for cur_epoch in range(100):
    optimizer.zero_grad()
    output = model(X)
    cur_loss = loss(output, y)
    cur_loss.backward()
    optimizer.step()
    cur_loss = loss(output, y)
    print("Epoch {0} Loss is {1}".format(cur_epoch, cur_loss))
    wandb.log({"loss": cur_loss})
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/e/ee2debd43d69f6037e1743011c16770f6169a03f.png" data-download-href="https://discuss.pytorch.org/uploads/default/ee2debd43d69f6037e1743011c16770f6169a03f" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/e/e/ee2debd43d69f6037e1743011c16770f6169a03f_2_690x352.png" alt="image" data-base62-sha1="xZ21t6iRccj94YgvRAZdpJZq2FF" width="690" height="352" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/e/e/ee2debd43d69f6037e1743011c16770f6169a03f_2_690x352.png, https://discuss.pytorch.org/uploads/default/optimized/3X/e/e/ee2debd43d69f6037e1743011c16770f6169a03f_2_1035x528.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/e/e/ee2debd43d69f6037e1743011c16770f6169a03f.png 2x" data-dominant-color="E8ECEF"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1166×595 28.7 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p><strong>MODIFIED</strong></p>
<pre><code class="lang-auto">import numpy as np
import time
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable

class random_model(nn.Module):
    def __init__(self, num_layers):
        super(random_model, self).__init__()
        self.model1 = nn.Sequential(nn.Linear(100, 20), nn.BatchNorm1d(20), nn.ReLU())
        self.model2 = []
        for i in range(num_layers):
            self.model2.append(nn.Sequential(nn.Linear(20, 20), nn.BatchNorm1d(20), nn.ReLU()))
        self.model2 = nn.Sequential(*self.model2)
        self.model3 = nn.Linear(20, 1)
    
    def forward(self, x):
        x = self.model1(x)
        #for cur_model in self.model2:
        #    x = cur_model(x)
        x = self.model2(x)
        x = self.model3(x)
        return x
    
model = random_model(10)
loss = torch.nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

X = torch.rand(100, 100)
y = torch.rand(100)

import wandb
wandb.init(project="random-project-custom-class", entity="")
wandb.config = {
  "learning_rate": 0.001,
  "epochs": 100,
  "batch_size": 100
}

# 3. Log gradients and model parameters
wandb.watch(model, log="gradients", log_freq=100)
print("DEF")
for cur_epoch in range(100):
    optimizer.zero_grad()
    output = model(X)
    cur_loss = loss(output, y)
    cur_loss.backward()
    optimizer.step()
    cur_loss = loss(output, y)
    print("Epoch {0} Loss is {1}".format(cur_epoch, cur_loss))
    wandb.log({"loss": cur_loss})
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/4/4/44e06588ccaa886a7261dcfdf6dccc2cc38b2080.png" data-download-href="https://discuss.pytorch.org/uploads/default/44e06588ccaa886a7261dcfdf6dccc2cc38b2080" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/4/4/44e06588ccaa886a7261dcfdf6dccc2cc38b2080_2_690x394.png" alt="image" data-base62-sha1="9PjeBn5hA00xm2CxkEgJBSDZgju" width="690" height="394" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/4/4/44e06588ccaa886a7261dcfdf6dccc2cc38b2080_2_690x394.png, https://discuss.pytorch.org/uploads/default/optimized/3X/4/4/44e06588ccaa886a7261dcfdf6dccc2cc38b2080_2_1035x591.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/4/4/44e06588ccaa886a7261dcfdf6dccc2cc38b2080.png 2x" data-dominant-color="F0F3F6"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1301×744 61.5 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

username: ColdFrenzy
cooked: <p><a class="mention" href="/u/anantguptadbl">@anantguptadbl</a><br>
Yes, I agree that the nn.Sequential model works like this and it is possible to view it on wandb, in fact in my second code snippet I am able to view the graphs of the gradient correctly.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/4/e4b490722f003030451a413a551c7bee3ddf39aa.png" data-download-href="https://discuss.pytorch.org/uploads/default/e4b490722f003030451a413a551c7bee3ddf39aa" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/e/4/e4b490722f003030451a413a551c7bee3ddf39aa_2_690x262.png" alt="image" data-base62-sha1="wDdIF5I3sn9z09kI9IFzM8OuHCq" width="690" height="262" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/e/4/e4b490722f003030451a413a551c7bee3ddf39aa_2_690x262.png, https://discuss.pytorch.org/uploads/default/optimized/3X/e/4/e4b490722f003030451a413a551c7bee3ddf39aa_2_1035x393.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/e/4/e4b490722f003030451a413a551c7bee3ddf39aa_2_1380x524.png 2x" data-dominant-color="F9FAFB"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1816×692 29.1 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>but the difference with your code is that in my case, at some point the network proceeds in parallel and not completely in sequence (the actor and critic subnets). In this case, if I pass the shared_layer output to the two subnets (which are two nn.Sequential modules) the loss values ​​become completely meaningless and I think it is due to the use of the nn.Sequential modules, since the rest of the program is exactly the same</p>

username: anantguptadbl
cooked: <aside class="quote no-group" data-username="ColdFrenzy" data-post="5" data-topic="141613">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/coldfrenzy/48/45754_2.png" class="avatar"> ColdFrenzy:</div>
<blockquote>
<p>is case, if I pass the shared_layer</p>
</blockquote>
</aside>
<p>Aah understood. Let me try that as well</p>