OPHoperHPO / image-background-remove-tool

✂️ Automated high-quality background removal framework for an image using neural networks. ✂️
https://carve.photos
Apache License 2.0
1.39k stars 272 forks source link

Question: Mismatch error for further pre-trained model #87

Closed dabernstein closed 2 years ago

dabernstein commented 2 years ago

I was training IS-Net pre-trained model further and got an mismatch error RuntimeError: Error(s) in loading state_dict for U2NET: Missing key(s) in state_dict: "outconv.weight", "outconv.bias". Unexpected key(s) in state_dict: "conv_in.weight", "conv_in.bias". size mismatch for stage1.rebnconvin.conv_s1.weight: copying a param with shape torch.Size([64, 64, 3, 3]) from checkpoint, the shape in current model is torch.Size([64, 3, 3, 3]).

My attempt for the training was so that the model would be able to handle higher res images. For my question it would be to solve this would I have to continue training the U2Net model instead or is it possible to train the IS-Net model to fix this?

OPHoperHPO commented 2 years ago

Hey @dabernstein

There is currently no out-of-the-box support for IS-Net. Before loading an IS-Net model into this tool, you need to implement an IS-Net wrapper that is compatible with the framework, then initialize it, load the model, and pass the wrapper class to the interface class.

Now I am preparing a major update for this framework, which will greatly improve the quality of background removal. You can wait for it, or add support for this neural network and send a Pull Request here :)