Memory allocation issue

ANKIT_SINGH_2128: <p>import torch<br>
import torch.nn as nn<br>
import torchvision.datasets as datasets<br>
import torchvision.transforms as transforms<br>
from torch.utils.data import DataLoader</p>
<p>class ChannelAttention(nn.Module):<br>
def <strong>init</strong>(self, in_channels, reduction_ratio=16):<br>
super(ChannelAttention, self).<strong>init</strong>()<br>
self.avg_pool = nn.AdaptiveAvgPool2d(1)<br>
self.max_pool = nn.AdaptiveMaxPool2d(1)<br>
self.fc1 = nn.Conv2d(in_channels, in_channels // reduction_ratio, kernel_size=1, stride=1, padding=0)<br>
self.relu = nn.ReLU()<br>
self.fc2 = nn.Conv2d(in_channels // reduction_ratio, in_channels, kernel_size=1, stride=1, padding=0)</p>
<pre><code>def forward(self, x):
    avg_out = self.fc2(self.relu(self.fc1(self.avg_pool(x))))
    max_out = self.fc2(self.relu(self.fc1(self.max_pool(x))))
    out = avg_out + max_out
    return torch.sigmoid(out) * x
</code></pre>
<p>class SpatialAttention(nn.Module):<br>
def <strong>init</strong>(self, kernel_size=7):<br>
super(SpatialAttention, self).<strong>init</strong>()<br>
assert kernel_size in (3, 7), “kernel size must be 3 or 7”<br>
padding = 3 if kernel_size == 7 else 1<br>
self.conv = nn.Conv2d(2, 1, kernel_size, padding=padding)</p>
<pre><code>def forward(self, x):
    avg_out = torch.mean(x, dim=1, keepdim=True)
    max_out, _ = torch.max(x, dim=1, keepdim=True)
    out = torch.cat([avg_out, max_out], dim=1)
    out = self.conv(out)
    return torch.sigmoid(out) * x
</code></pre>
<p>class DualAttentionModule(nn.Module):<br>
def <strong>init</strong>(self, in_channels):<br>
super(DualAttentionModule, self).<strong>init</strong>()<br>
self.channel_attention = ChannelAttention(in_channels)<br>
self.spatial_attention = SpatialAttention()</p>
<pre><code>def forward(self, x):
    out = self.channel_attention(x) * x
    out = self.spatial_attention(out) * out
    return out
</code></pre>
<p>class DualAttentionNetwork(nn.Module):<br>
def <strong>init</strong>(self, num_classes=10):<br>
super(DualAttentionNetwork, self).<strong>init</strong>()<br>
self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1)<br>
self.relu = nn.ReLU()<br>
self.dual_attention1 = DualAttentionModule(64)<br>
self.conv2 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)<br>
self.pool = nn.MaxPool2d(kernel_size=2, stride=2)<br>
self.dual_attention2 = DualAttentionModule(128)<br>
self.fc1 = nn.Linear(128 * 4 * 4, 1024)  # Adjusted the input size for fc1<br>
self.fc2 = nn.Linear(1024, num_classes)</p>
<pre><code>def forward(self, x):
    out = self.conv1(x)
    out = self.relu(out)
    out = self.dual_attention1(out)
    out = self.conv2(out)
    out = self.relu(out)
    out = self.pool(out)
    out = self.dual_attention2(out)
    out = out.view(out.size(0), -1)
    out = self.fc1(out)
    out = self.relu(out)
    out = self.fc2(out)
    return out
</code></pre>
<h1><a name="set-device-gpu-or-cpu-1" class="anchor" href="#set-device-gpu-or-cpu-1"></a>Set device (GPU or CPU)</h1>
<p>device = torch.device(“cuda” if torch.cuda.is_available() else “cpu”)</p>
<h1><a name="define-hyperparameters-2" class="anchor" href="#define-hyperparameters-2"></a>Define hyperparameters</h1>
<p>num_epochs = 10<br>
batch_size = 2<br>
learning_rate = 0.001</p>
<h1><a name="load-the-dataset-and-apply-transformations-3" class="anchor" href="#load-the-dataset-and-apply-transformations-3"></a>Load the dataset and apply transformations</h1>
<p>transform = transforms.Compose([<br>
transforms.Resize((224, 224)),<br>
transforms.ToTensor(),<br>
transforms.Normalize((0.5,), (0.5,))<br>
])</p>
<p>train_dataset = datasets.ImageFolder(root=r’C:\Users\tmlan\OneDrive\Desktop\MODIFIED\UAVWeedSegmentation-main\training-data\trainval’, transform=transform)<br>
train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)</p>
<h1><a name="initialize-the-dual-attention-network-4" class="anchor" href="#initialize-the-dual-attention-network-4"></a>Initialize the Dual Attention Network</h1>
<p>model = DualAttentionNetwork().to(device)</p>
<h1><a name="define-loss-function-and-optimizer-5" class="anchor" href="#define-loss-function-and-optimizer-5"></a>Define loss function and optimizer</h1>
<p>criterion = nn.CrossEntropyLoss()<br>
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)</p>
<h1><a name="training-loop-6" class="anchor" href="#training-loop-6"></a>Training loop</h1>
<p>total_steps = len(train_loader)<br>
for epoch in range(num_epochs):<br>
for i, (images, labels) in enumerate(train_loader):<br>
images = images.to(device)<br>
labels = labels.to(device)</p>
<pre><code>    # Forward pass
    outputs = model(images)
    loss = criterion(outputs, labels)

    # Backward pass and optimize
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    # Print training progress
    if (i + 1) % 100 == 0:
        print(f"Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{total_steps}], Loss: {loss.item():.4f}")
</code></pre>
<h1><a name="save-the-trained-model-7" class="anchor" href="#save-the-trained-model-7"></a>Save the trained model</h1>
<p>torch.save(model.state_dict(), “dual_attention_model.pt”)<br>
I am getting error as :<br>
RuntimeError: mat1 and mat2 shapes cannot be multiplied (2x1605632 and 2048x1024)</p>

ptrblck: <aside class="quote no-group" data-username="ANKIT_SINGH_2128" data-post="1" data-topic="183963">
<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/ankit_singh_2128/48/61680_2.png" class="avatar"> ANKIT_SINGH_2128:</div>
<blockquote>
<p>I am getting error as :<br>
RuntimeError: mat1 and mat2 shapes cannot be multiplied (2x1605632 and 2048x1024)</p>
</blockquote>
</aside>
<p>This shape mismatch is not a memory allocation issue and is raised in <code>self.fc1</code>.<br>
Either set its <code>in_features</code> to <code>1605632</code> or reduce the size of the incoming activation to create <code>2048</code> features.</p>

ANKIT_SINGH_2128: <p>I am not able to understand could you pls tell which part of the code needs to be modified so that it doesnot cause any issue , please…</p>

ptrblck: <p>Replace:</p>
<pre><code class="lang-python">self.fc1 = nn.Linear(128 * 4 * 4, 1024)
</code></pre>
<p>with:</p>
<pre><code class="lang-python">self.fc1 = nn.Linear(1605632, 1024)
</code></pre>

ANKIT_SINGH_2128: <p>Again getting this error.When replaced with above.<br>
RuntimeError: [enforce fail at C:\actions-runner_work\pytorch\pytorch\builder\windows\pytorch\c10\core\impl\alloc_cpu.cpp:72] data. DefaultCPUAllocator: not enough memory: you tried to allocate 6576668672 bytes.</p>

ptrblck: <p>In that case reduce the input size or change the model architecture so that the activation input to this linear layer will be smaller, since you are running out of memory.</p>