RNN generator weights not being updated

username: marco_07
cooked: <p>Hello there,<br>
I’m currently working on a GAN that has 2 paired RNNs as generator (rnn and output in the following code).<br>
My problem is that the weights of the generator are not being updated.<br>
The following code is used to perform a sigle forward step of the 2 RNNs to generate a single observation.<br>
The main idea is that I have the tensor y_pred_long that is filled step after step with the predictions of the 2 RNNs.<br>
Then this structure (y_pred_long) is passed to a decoding function that retrieves the desidered input for the discriminator.<br>
I not beeing able to understand whether/where the computational graph breaks down, any suggestion is very appreciated!<br>
If anything from the following code is not clear and additional information is needed please let me know! (first time posting here)</p>
<pre><code class="lang-python">
def gen_rnn_epoch(epoch, rnn, output, device):

    rnn.hidden = rnn.init_hidden(batch_size=1)  # this is the source of noise for the GAN generator

    # empty tensor that will be filled at each step with predictions from both RNNs
    y_pred_long = torch.zeros(1, 9, 8 * 5 + 4).to(device)
    # the first input of RNN_1 is a token
    x_step = torch.ones(1, 1, 8 * 5 + 4).to(device)

    # sequential generation of RNN_1
    for i in range(9):

        h, node_prediction = rnn(x_step)
        node_prediction = F.gumbel_softmax(node_prediction, dim=2, hard=True)

        # init RNN_2: context init
        hidden_null = torch.zeros(rnn.num_layers - 1, h.size(0), h.size(2)).to(device)
        output.hidden = torch.cat((h.permute(1, 0, 2), hidden_null), dim=0)

        # reset token and fill with newest prediction
        x_step = torch.zeros(1, 1, 8 * 5 + 4).to(device)
        x_step[:, :, :4] = node_prediction

	    # the first input of RNN_2 is a token
        output_x_step = torch.ones(1, 1, 5).to(device)
	
		# sequential generation of RNN_2
        edge_rnn_step = 0
        for j in range(min(8, i + 1)):
            output_y_pred_step = output(output_x_step)
            output_x_step = F.gumbel_softmax(output_y_pred_step, dim=2, hard=True)
            x_step[:, :, 4 * j + 4 + j: 4 * (j + 1) + 4 + (j + 1)] = output_x_step
            output.hidden = output.hidden.to(device)
            edge_rnn_step = j

        y_pred_long[:, i:i + 1, :] = x_step

        # stop criteria (EOS met)
        node_to_break, edges_to_break = torch.split(x_step, [4, 40], dim=2)
        edges_to_break_temp = torch.reshape(edges_to_break, (edges_to_break.shape[0], 8, 5))
        edges_to_break_uptillnow = edges_to_break_temp[0, :edge_rnn_step + 1, :]
        break_ = True
        for row in edges_to_break_uptillnow:
            if torch.argmax(row).item() != 0:
                break_ = False
            if edge_rnn_step == 0:
                break_ = False

        if break_:
            break

        rnn.hidden = rnn.hidden.to(device)

    x, edg_idx, edg_attr = decode_adj_generation(y_pred_long[0], device)

    data = Data(x=x, edge_index=edg_idx, edge_attr=edg_attr)

    return data

def decode_adj_generation(encoded_adj, device):

    max_prev_TEST = encoded_adj.shape[0] - 1  
    node_x, edges = torch.split(encoded_adj, [4, 40], dim=1)
    edges = torch.reshape(edges, (encoded_adj.shape[0], max_prev_TEST, 5))

    # drop all the rows for which we didn't have sampled a node
    idxs_to_keep = []
    for row in range(node_x.shape[0]):
        if (torch.sum(node_x[row]).item()) != 0:
            idxs_to_keep.append(row)
        else:
            break

    idxs_to_keep = torch.LongTensor(idxs_to_keep).to(device)
    node_x = node_x.index_select(0, idxs_to_keep)
    edges = edges.index_select(0, idxs_to_keep)

    # un-FLIP the edges
    for i in range(edges.shape[0]):
        edges[i, :i + 1, :] = torch.flip(edges[i, :i + 1, :], [0]) 

    # below for edge index creation, this step CAN be NON-differentiable
    first_dim_idx = []
    second_dim_idx = []

    for i in range(edges.shape[0]):
        for j in range(i):
            if torch.sum(edges[i - 1, j, :]) != 0:
                if np.nonzero(edges[i - 1, j, :])[0].item() != 0:
                    first_dim_idx.append(i)
                    second_dim_idx.append(j)

    first_dim_idx = torch.tensor(torch.LongTensor(first_dim_idx)).to(device)
    second_dim_idx = torch.tensor(torch.LongTensor(second_dim_idx)).to(device)
    row1 = torch.unsqueeze(torch.cat((first_dim_idx, second_dim_idx), 0), 0)
    row2 = torch.unsqueeze(torch.cat((second_dim_idx, first_dim_idx), 0), 0)
    edge_idx = torch.cat((row1, row2), 0)
    half_edge_idx = torch.cat((torch.unsqueeze(second_dim_idx, 0), torch.unsqueeze(first_dim_idx - 1, 0)), 0)

    # construction of edge_attr matrix, this has to be differentiable
    edge_attr = None
    step = 0
    for i in range(first_dim_idx.shape[0]):
        if step == 0:
            edge_attr = torch.unsqueeze((edges[half_edge_idx[1][i].item(), half_edge_idx[0][i].item(), 1:]), 0)
            step += 1
        else:
            edge_attr = torch.cat(
                (edge_attr, torch.unsqueeze(edges[half_edge_idx[1][i].item(), half_edge_idx[0][i].item(), 1:], 0)), dim=0)

    if edge_attr is not None:
        edge_attr = torch.cat((edge_attr, edge_attr), 0)
    else:
        edge_attr = torch.empty(0, 4)

    return node_x.to(torch.float32).to(device), edge_idx.to(torch.long).to(device), edge_attr.to(torch.float32).to(
            device)
</code></pre>