问根据层次结构将pytorch模型参数导出到不同的文件中

EN

class ConvBlock(nn.Module):
    def __init__(self, C_in, C_out, kernel, pool):
        super().__init__()
        self.conv = nn.Conv2d(C_in, C_out, kernel)
        self.relu = nn.ReLU(inplace = True)
        self.pool = nn.MaxPool2d(2,2) if pool else nn.Identity()
        
    def forward(self, input):
        out = self.conv(input)
        out = self.relu(out)
        out = self.pool(out)
        return out
    
class LeNet5(nn.Module):
    def __init__(self):
        super().__init__()
        self.block1 = ConvBlock(1, 6, 5, pool = True)
        self.block2 = ConvBlock(6, 16, 5, pool = True)
        self.block3 = ConvBlock(16, 120, 5, pool = False)
        self.fc = nn.Sequential(
            nn.Linear(120, 84),
            nn.ReLU(inplace = True),
            nn.Linear(84, 10)
        )
        
    def forward(self, input):
        out = self.block1(input)
        out = self.block2(out)
        out = self.block3(out)
        out = out.view(-1,120)
        out = self.fc(out)
        return out

net = LeNet5()

basedir = 'lenet_params'

for name, param in net.named_parameters():
    name = name.split('.')
    out_dir, filename = os.path.join(basedir, *name[:-1]), name[-1]+'.pth'
    out_path = os.path.join(out_dir, filename)

    if not os.path.exists(out_dir):
        os.makedirs(out_dir, exist_ok=True)
    torch.save(param, out_path)

lenet_params
|---block1
|   |---conv
|   |   |---weight.pth
|   |   |---bias.pth
|---block2
|   |---conv
|   |   |---weight.pth
|   |   |---bias.pth
|---block3
|   |---conv
|   |   |---weight.pth
|   |   |---bias.pth
|---fc
|   |---0
|   |   |---weight.pth
|   |   |---bias.pth
|   |---2
|   |   |---weight.pth
|   |   |---bias.pth