How can i remove layers from torch_model?

1

I’m trying to remove layers from this model and at the same time keep the “TorchModel”, how can I do that? I tried using .children, but a problem occurs in another line of code.

image
image783×489 19.2 KB

2

I don’t know the direct method to modify the name of class. Instead, use .named_children()

class TorchModel(nn.Module):
    def __init__(self, torch_model):
        super(TorchModel, self).__init__()
        for n, m in torch_model.named_children():
            self.__setattr__(n, m)

    def forward(self, x):
        # whatever you want

my_model = TorchModel(torch_model)