How to disactivate a layer at evaluation

vision
1

Hi,

I am training a ResNet18 with a graph neural network on top of it, I want to discard the GAT layer at evaluation, can someone help me with this?

def forward(self, x, lin=0, lout=5):
    out = x
    if lin < 1 and lout > -1:
        out = self.conv1(out)
        out = self.bn1(out)
        out = F.relu(out)
    if lin < 2 and lout > 0:
        out = self.layer1(out)
    if lin < 3 and lout > 1:
        out = self.layer2(out)
    if lin < 4 and lout > 2:
        out = self.layer3(out)
    if lin < 5 and lout > 3:
        
         #This part
        out = out.view(out.size(0), -1)
        edges = self.topology(self.knn(out),512)
        graph_data = (out,edges,edges)
        out,edge,attention = self.gat_net(graph_data)
        out = out.squeeze()
        out = torch.reshape(out,(512,256,8,8))
        # to here,  I don't want it to perform at evaluation
        out = self.layer4(out)
    if lout > 4:
        #out = F.avg_pool2d(out,2)
        #res = out
        out = F.adaptive_avg_pool2d(out, (1, 1))
        out = out.view(out.size(0), -1)
        out = self.linear(out)

    return out

Thanks

2

You could add a condition using the internal self.training argument:

        out = torch.reshape(out,(512,256,8,8))
        # to here,  I don't want it to perform at evaluation
        if self.training: # will only be called during training
            out = self.layer4(out)

The self.training argument is set to True or False by calling model.train() and model.eval(), respectively.