Updating an instance attribute within a custom Conv2D

Hello,

I am trying to modify a custom Conv2D’s custom attribute, but failing as of now.

For example, in the following custom Conv2D which inherits from PyTorch’s nn.Conv2d, I have a ‘count’ attribute that should get counted up each time the forward function in an instantiated Conv2d object is called. However, the value of count resets to 0 (i.e. the printed output is always 0) when the forward function is invoked every time. Can you please point out what I am doing wrong here?

Thanks in advance!

Code snippet:
class Conv2d(nn.Conv2d):

def __init__(self, in_channels, out_channels, kernel_size, stride=1,
             padding=0, dilation=1, groups=1, bias=True):
    super(Conv2d, self).__init__(in_channels, out_channels, kernel_size, stride,
             padding, dilation, groups, bias)

    self.count = 0

def forward(self, x):
    weight = self.weight
    print(self.count)  # always prints 0
    self.count += 1
    return F.conv2d(x, weight, self.bias, self.stride,
                    self.padding, self.dilation, self.groups)

I run your code and it prints 0, 1, 2, 3, 4, 5, 6, 7, 8, 9.

import torch.nn as nn
import torch.nn.functional as F

class Conv2d(nn.Conv2d):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1,
             padding=0, dilation=1, groups=1, bias=True):
        super(Conv2d, self).__init__(in_channels, out_channels, kernel_size, stride,
             padding, dilation, groups, bias)

        self.count = 0

    def forward(self, x):
        weight = self.weight
        print(self.count)  # always prints 0
        self.count += 1
        return F.conv2d(x, weight, self.bias, self.stride,
                    self.padding, self.dilation, self.groups)
        
net = Conv2d(3, 3, 3)
for _ in range(10):
    x = torch.rand(3, 3, 3, 3)
    net(x)
1 Like

Thank you for looking into this, Eta_C.

You are right - it counts up with CPU models, but this seems to break for the same model with DataParallel and CUDA. Please see the code snippet below.

import torch.nn
import torch.nn as nn
import torch.nn.functional as F

class Conv2d(nn.Conv2d):
def init(self, in_channels, out_channels, kernel_size, stride=1,
padding=0, dilation=1, groups=1, bias=True):
super(Conv2d, self).init(in_channels, out_channels, kernel_size, stride,
padding, dilation, groups, bias)

    self.count = 0

def forward(self, x):
    weight = self.weight
    print(self.count)  # always prints 0
    self.count += 1
    return F.conv2d(x, weight, self.bias, self.stride,
                self.padding, self.dilation, self.groups)

net = Conv2d(3, 3, 3)
net = torch.nn.DataParallel(net).cuda()
for _ in range(10):
x = torch.rand(3, 3, 3, 3)
x = x.cuda()
net(x)

You are right. I found that I have to use multiprocessing.

The following code will work but,

  • it is not elegant and generic enough
  • If I have 2 GPUs, it will output 1 to 20
  • I don’t know if there are any other side effects
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import multiprocessing as mp

class Conv2d(nn.Conv2d):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.cnt = mp.Value("d", 0)
    
    def forward(self, x):
        self.cnt.value += 1
        return super().forward(x)

net = Conv2d(3, 3, 3)
net = nn.DataParallel(net).cuda()
for _ in range(10):
    x = torch.rand(3, 3, 3, 3)
    x = x.cuda()
    net(x)
1 Like

Thank you very much for looking into this. The solution you mentioned works for me. I do hope though that later PyTorch builds will have a more generic solution to this problem.

Thank you once again for helping out.