Recompute convolution without parameter update

[jakubMitura14]

Hello, I have unet like architecture, so when x is input, pseudocode

Out1= conv1(x)
Out2 =conv2(out1) 
Out3= deconv1(out2)+ out1
Return deconv2(out3) +x

Now, the problem is that I am memory constrained, so instead of keeping all in memory, I would do remat and sequential, so line 3 will become

Deconv1(out2)+conv1(x)

However, for this case, I need to recompute
conv1(x) but avoid any parameter learning that was already done in line 1

I know that one can freeze parameters using optax, but as far as I get it, it would freeze all conv1 cases, so both first as well as second invocation, and I want to block only the second one.

Is there a way to selectively block parameter learning?

[cgarciae]

Hey @jakubMitura14, take a look at nn.map_variables, you can probably use a similar pattern as the example but apply_mask should instead call stop_gradient on the desired parameters.