PyTorch summary 和 冻结网络层

本文是最近研究Prompt Tuning的时候，想使用 torchsummary 来打印一下模型及可训练参数。

于是有了下面的这个例子。

import torch
import torch.nn as nn
from torchsummary import summary

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.embedding = nn.Embedding(50,32)
    def forward(self, input):
        # print(input)
        drop_1 = self.embedding(input)
        return drop_1
        

class MyModel1(nn.Module):
    def __init__(self):
        super(MyModel1, self).__init__()
        self.model = MyModel()
        """"""
        for param in self.model.parameters():
        	param.requires_grad = False
        
        self.embedding = nn.Embedding(50,32)
	
    def forward(self, input):
        out0 = self.model(input)
        out1 = self.embedding(input)
        a = out0[:, :25].clone()
        b = out0[:, 25:].clone()
        bools1 = torch.cat((a.fill_(False), b.fill_(True)), axis = 1)
        bools2 = torch.cat((a.fill_(Tru