Running inference on mobile about ace

[daniel-sudz]

it looks this issue would be a problem with making the model scriptable: pytorch/pytorch#36061
https://github.com/nianticlabs/ace/blob/main/ace_network.py#L91

one alternative would be to do something like this:

        self.res_blocks = nn.ModuleList(
            nn.ModuleList(([
                nn.Conv2d(self.head_channels, self.head_channels, 1, 1, 0),
                nn.Conv2d(self.head_channels, self.head_channels, 1, 1, 0),
                nn.Conv2d(self.head_channels, self.head_channels, 1, 1, 0),
            ])) for block in range(num_head_blocks)
        )

though this would end up breaking this import:

ace/ace_network.py

Line 212 in 6b2a3bf

pattern = re.compile(r"^heads\.\d+c0\.weight$")

[daniel-sudz]

@tcavallari I would really love to get inference to run on mobile if it's possible, but it appears as though the pretrained weights must be retrained as per this issue (I'm assuming that there are currently no plans to make the code for the pretrained weights available):

pytorch/pytorch#36061

To make these projects scriptable, we must replace list of modules with nn.ModuleList which will make the pretrained models invalid, since the replacement will change the keys of state_dict, and may increase the number of model parameters.

[tcavallari]

Hello!

To make these projects scriptable, we must replace list of modules with nn.ModuleList which will make the pretrained models invalid, since the replacement will change the keys of state_dict, and may increase the number of model parameters.

I don't think it's necessary to retrain the encoder. You can create a scriptable architecture using nn.ModuleList and inserting the same convolution layers (keeping the same parameters, i.e. kernel size, stride, padding, etc...) inside it. The number of model parameters shouldn't change at all.

Then it should just be a matter of replacing the weights in the new state_dict with the weights in the pretrained one we provide with the repo. You just have to figure out which keys in the new dict map to which keys in the old dictionary, but it should be relatively easy.

[daniel-sudz]

Hello!

To make these projects scriptable, we must replace list of modules with nn.ModuleList which will make the pretrained models invalid, since the replacement will change the keys of state_dict, and may increase the number of model parameters.

I don't think it's necessary to retrain the encoder. You can create a scriptable architecture using nn.ModuleList and inserting the same convolution layers (keeping the same parameters, i.e. kernel size, stride, padding, etc...) inside it. The number of model parameters shouldn't change at all.

Then it should just be a matter of replacing the weights in the new state_dict with the weights in the pretrained one we provide with the repo. You just have to figure out which keys in the new dict map to which keys in the old dictionary, but it should be relatively easy.

Ok awesome I think I managed to figure it out. I believe this change should be completely seamless in terms of back-compat so do you think we could merge this into main? I made a PR here in #22. Thank you so much for the guidance