I have some questions about the usage of local variables in the forward function. Say I have the following model:
class Model(nn.Module):
def __init__(self):
self.conv1 = nn.Conv2d(3,64,1)
self.bn1 = nn.BatchNorm2d(64)
self.conv2 = nn.Conv2d(64,128,1)
self.bn2 = nn.BatchNorm2d(128)
self.relu = nn.ReLU(inplace=True)
I can write the forward function in 2 ways:
(1)
def forward(self, x):
x = self.relu(self.bn1(self.conv1(x)))
x = self.relu(self.bn2(self.conv2(x)))
return x
(2)
def forward(self, x):
x = self.relu(self.bn2(self.conv2(self.relu(self.bn1(self.conv1(x))))))
return x
Does (2) use less memory than (1)?
Also, when relu is inplace, I get an error if I do:
x = conv1(x)
x = bn1(x)
x = relu(x)
However, it is fine if I do:
x = relu(bn1(conv1(x)))
Is it because relu is inplace I can’t assign a local variable to its result?