pytorch gets the middle layer output result of the model

In the inference stage, the entire model will be loaded on the GPU for end-to-end calculations, usually only giving you a final result.

If you want to obtain the middle layer output of the model, you need to calibrate the target layer position (return by forward) before calculation , or cut off the model at that layer (as a small model) and output . This article introduces two methods to obtain the output results of the middle layer of the model:

1. Calibrate the position of the target layer

def forward(self, x):
    layer1_out = self.layer1(x)
    layer2_out = self.layer2(layer1_out)
    out = self.final_layer(layer2_out)
    return out #, layer2_out

If you want to see the output result of the second layer, just return layer2_out when returning.

2. Cut off the model at a certain level

class MyModel:
    def __init__(self, xxx):
        self.backbone = net1()
        self.head = net2()
    def forward(self, x):
        out = self.head(net1(x))
        return out

if __name__=='__main__':
    model = MyModel()
    x = cv2.imread('xxx')
    mid_out = model.backbone(x)
    print(mid_out.size())

 

Guess you like

Origin blog.csdn.net/leviopku/article/details/108474310