Heap Corruption error for TGNN

Matthew_Giacovelli: <p>Hello,<br>
I’m getting the following error Process finished with exit code -1073740940 (0xC0000374).<br>
When I run my training loop for the TGNN. I get this error at various points, it may be during backpropagation, it may be during the forward pass, it might not make it though the first forward pass. Currently environment has multiple agents where each node is an agent with 10 features and the edges have an attribute which is the distance between the agents. When I vary the number of agents, usually the process will go farther when we have more agents. I’m not sure if C++ is the right category, but I believe the issue with the memory is something happening with the underlying C++.</p>
<p>Below I am attaching my architecture and the training loop. My environment is a custom gym environment. Small note, the dummy_loss is just an artifact from trying to get backprop() to work. I thought maybe the reason sometimes I would get the corruption error was I was not using all the predicted outputs from the model. I also tried adding more edges because in situations where my graph was small maybe there was a problem with the message passing. However neither of these would explain why sometimes I don’t make it past the first layer of the forward pass.</p>
<p>import torch<br>
import torch.nn.functional as F<br>
from torch_geometric.nn import GCNConv<br>
from torch_geometric.data import Data</p>
<p>class TemporalGNN(torch.nn.Module):<br>
def <strong>init</strong>(self):<br>
super(TemporalGNN, self).<strong>init</strong>()<br>
self.conv1 = GCNConv(in_channels=10, out_channels=32)<br>
self.conv2 = GCNConv(in_channels=32, out_channels=16)<br>
self.conv3 = GCNConv(in_channels=16, out_channels=8)<br>
self.conv4 = GCNConv(in_channels=8, out_channels=4)<br>
self.conv5 = GCNConv(in_channels=4, out_channels=2)</p>
<pre><code>def forward(self, data):
    x, edge_index = data.x, data.edge_index
    x = self.conv1(x, edge_index)
    print("Conv 1")
    x = F.relu(x)
    print("RELU 1")
    x = self.conv2(x, edge_index)
    print("Conv 2")
    x = F.relu(x)
    print("RELU 2")
    x = self.conv3(x, edge_index)
    print("Conv 3")
    x = F.relu(x)
    print("RELU 3")
    x = self.conv4(x, edge_index)
    print("Conv 4")
    x = F.relu(x)
    print("RELU 4")
    x = self.conv5(x, edge_index)
    print("Conv 5")
    x_norm = torch.norm(x, p=2, dim=1, keepdim=True)

    # Prevent division by zero by adding a small epsilon
    epsilon = 1e-8
    x_norm = x_norm + epsilon
    x = x / x_norm

    """
    # Only keep predictions for pursuers
    predicted_velocities = x[pursuer_indices]
    """

    return x
</code></pre>
<p>import Environment_FTP as e<br>
import pygame<br>
import torch<br>
import misc as m<br>
import define_TGN as tgn<br>
import sys<br>
import matplotlib.pyplot as plt<br>
import psutil<br>
import os<br>
os.environ[“CUDA_LAUNCH_BLOCKING”] = “1”</p>
<p>def train_and_evaluate(_screen, _robots, _adversaries, _episodes, _time_steps, _epochs, _learning_rate):<br>
torch.set_printoptions(sci_mode=False, precision=4)</p>
<pre><code># Set up model
model = tgn.TemporalGNN()

optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

overseer = e.Overseer(_screen, _robots, 0, _adversaries, _time_steps, 15)

wins = 0
losses = []

# loop for what we do each episode
for episode in range(_episodes):
    # reset the environment
    overseer.reset()
    self_dict = overseer.o

    # keeps track of the losses from each timestep
    total_loss = torch.tensor(0.0, requires_grad=True)

    # Make attributes a tensor
    for agent in self_dict["all_agents"]:
        self_dict["all_agents"][agent].attributes["x"] = torch.tensor(
            self_dict["all_agents"][agent].attributes["x"], dtype=torch.float32, requires_grad=True)

        self_dict["all_agents"][agent].attributes["y"] = torch.tensor(
            self_dict["all_agents"][agent].attributes["y"], dtype=torch.float32, requires_grad=True)

        self_dict["all_agents"][agent].attributes["v_x"] = torch.tensor(
            self_dict["all_agents"][agent].attributes["v_x"], dtype=torch.float32, requires_grad=True)

        self_dict["all_agents"][agent].attributes["v_y"] = torch.tensor(
            self_dict["all_agents"][agent].attributes["v_y"], dtype=torch.float32, requires_grad=True)

        self_dict["all_agents"][agent].attributes["c"] = torch.tensor(
            self_dict["all_agents"][agent].attributes["c"], dtype=torch.float32, requires_grad=True)

        self_dict["all_agents"][agent].attributes["status"] = torch.tensor(
            self_dict["all_agents"][agent].attributes["status"], dtype=torch.float32)

        self_dict["all_agents"][agent].attributes["r"] = torch.tensor(
            self_dict["all_agents"][agent].attributes["r"], dtype=torch.float32)

        self_dict["all_agents"][agent].attributes["W"] = torch.tensor(
            self_dict["all_agents"][agent].attributes["W"], dtype=torch.float32)

        self_dict["all_agents"][agent].attributes["deaths"] = torch.tensor(
            self_dict["all_agents"][agent].attributes["deaths"], dtype=torch.float32)

    # Gets data for TGN all in tensors
    data = m.get_data(self_dict)

    for t in range(_time_steps):

        # Now we normalize our features
        data, scalar = m.normalize_data(data)

        # returns list of pursuers
        pursuer_indices = m.get_pursuer_indices(self_dict)

        # TGN needs the agents as numbers
        n2a = m.numbers_as_agents(self_dict)

        # In train() we have dropout and batch
        model.train()

        robot_mask = []
        adversary_mask = []

        for agent_id, agent in self_dict["all_agents"].items():
            if agent_id.startswith("C") and agent.attributes["status"] == 1:
                robot_mask.append(1)  # Robot is active
                adversary_mask.append(0)  # Not an adversary
            else:
                robot_mask.append(0)  # Not a robot (or inactive)
                adversary_mask.append(1)  # This is an adversary

        robot_mask = torch.tensor(robot_mask, dtype=torch.float32)  # Shape: (num_agents,)
        adversary_mask = torch.tensor(adversary_mask, dtype=torch.float32)  # Shape: (num_agents,)

        # Get current predicted velocities
        predicted_velocities = model(data)

        robot_velocities = predicted_velocities * robot_mask.unsqueeze(1)  # Keep only robot velocities
        adversary_velocities = predicted_velocities * adversary_mask.unsqueeze(1)  # Keep only adversary velocities

        # Get number of adversaries before update
        old_counts = m.number_of_CDE(self_dict)[2]

        # Apply robot velocities from TGNN predictions, only for active robots (mask == 1)
        for idx, (key, agent) in enumerate(self_dict["all_agents"].items()):
            if robot_mask[idx] == 1:  # Only update if it's an active robot
                vx, vy = robot_velocities[idx]

                # Use the agent's ID as the key in the update_velocity dictionary
                self_dict["update_velocity"][key] = (vx, vy)

        # Add a small dummy loss to preserve adversary predictions in the graph
        dummy_loss = torch.sum(adversary_velocities ** 2) * 1e-8

        # Apply new actions
        obs, reward, terminateds, truncateds, info = overseer.step([total_loss, optimizer, model])

        # Get TI post update
        TI = m.get_total_influence(self_dict)

        # Get adversaries post update
        new_counts = m.number_of_CDE(self_dict)[2]
        # get_new_counts = new_counts.clone().detach().requires_grad_(True)

        # See if # adversaries changed
        del_a = old_counts - new_counts

        # calculate loss
        loss_t = -m.custom_reward(self_dict, TI, del_a) - dummy_loss
        print("loss at timestep", str(t + 1) + ":", loss_t)

        # Add loss to total for episode
        total_loss = total_loss + loss_t
        print("Total loss:", total_loss)
        print()

        if terminateds or truncateds:
            print("===============================================")
            print(f"Episode {episode + 1} ended with reward: {reward}")
            wins += 1 if info["val"] == 1 else 0
            break

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()

        data = m.un_normalize(data, scalar)

        m.update_graph(data, self_dict)

        # CPU usage
        cpu_percent = psutil.cpu_percent()
        print(f"CPU Usage: {cpu_percent}%")

        # RAM usage
        ram = psutil.virtual_memory()
        ram_used = ram.used / (1024 ** 3)  # Convert bytes to GB
        ram_total = ram.total / (1024 ** 3)  # Convert bytes to GB
        print(f"RAM Usage: {ram_used:.2f} GB / {ram_total:.2f} GB")

    # Backpropagation loss
    print(total_loss.requires_grad)
    total_loss.backward()
    optimizer.step()

    print()
    print("Ep grad")
    for name, param in model.named_parameters():
        if param.grad is not None:
            print(f"Gradient for {name}: {param.grad}")
        else:
            print(f"No gradient for {name}")

    print()
    print("Total Loss", total_loss.item())
    print()
    optimizer.zero_grad()
    for name, param in model.named_parameters():
        if param.grad is not None:
            print(f"Gradient for {name}: {param.grad}")
        else:
            print(f"No gradient for {name}")

# Plot episode losses
plt.plot(range(len(losses)), losses)
plt.show()
</code></pre>
<p>def main():<br>
pygame.init()<br>
screen = pygame.display.set_mode([1000, 1000])</p>
<pre><code>robots = 2
adversaries = 1
episodes = 100
time_steps = 500
epochs = 200
alpha = 0.01

train_and_evaluate(screen, robots, adversaries, episodes, time_steps, epochs, alpha)
</code></pre>
<p>main()</p>
<p>Thank you for your help!</p>