Is it possible to optimize only selected part of the tensor in each iteration?

Konsthr: <p>Hello!<br>
Is there a way to optimiize only a certain part of a tensor?<br>
I am truing to use Pythorch as a backend for ptychographycal engine, which requires to optimize  large [2^12,2^12,2] tensor O. But I can formulate loss function only for set of the overlaping slices of this big tensor O[up:down,left:right,:] so I need to calculate updates for each of these slices consistently. Currently, I just iterate through list of these slices and use optimizer with respect to the whole big tensor which  require to calclate gradient for the whole O each step which takes most of the time, despite I know that only small part of it O[up:down,left:right,:]  influence on the loss function and should be updated this step.</p>
<p>So if there any way to turn on gradient calculation only  for some part of O  each step, or are there any other ways to solve this problem?</p>

ptrblck: <p>How are you calculating the gradients at the moment?<br>
Could you post some dummy code, as I’m not sure how the loss is being calculated at the moment?</p>

Konsthr: <p>Currently I utilize the following strategy:<br>
Initially I have apriory known <code>Propagation_model</code>,<code>Measured_th</code> and <code>Slice_coordinates</code> where:<br>
<code>Propagation_model</code>  - model object, which <code>forward</code> method transform my guessed input data into guessed approximation of measured data. This method may include several FFT and multiplications with different kernels, but it is known percicely  in advance and does not require any optimization, so   <code>Propagation_model</code>  does not require gradient.<br>
<code>Measured_th</code> - <code>[n,1024,1024,2]</code> complex tensor which represents n measurments obtained during experiment and also requires no optimization.<br>
<code>Slice_coordinates</code>  - list of tuples of coordinates <code>[(up,down,left,right)]</code> - which allows to get part of the (<code>Sample_th</code> - tensor which represent  reconstructed  object ) corresponding to exact measurment  from  <code>Measured_th</code></p>
<p>During the optimizational procedure I trying to obtain two complex  tensors <code>Sample_th</code> <code>[4096,4096,2]</code> and <code>Probe_th</code> <code>[1024,1024,2]</code>, both of them requre to be optimized and require gradient.</p>
<p>My loss defined as sumed squared difference between certain slice of  ‘Measured_data’ and result of propagation of  certain slice of  <code>Sample_th</code> and <code>Probe_th</code></p>
<p>Before the begining of optimization loop I create set of slices <code>Slices_th</code> from <code>Sample_th</code>, since only a part <code>[1024,1024,2]</code> of  <code>Sample_th</code> with <code>(up,down,left,right)</code> coordinates takes part in production of certain <code>Measured_th</code> slice:<br>
<code>Slices_th= [] for num in range(len(Measured_th)): Slices_th.append((Sample_th[Borders.up[num]: Borders.down[num], Borders.left[num]: Borders.right[num],:]))</code></p>
<p>all members of <code>Slices_th</code> represenc slices taken from <code>Sample_th</code> and partially overlap each other ( certain part of  <code>Sample_th</code> belong to multiple memders of  <code>Slices_th</code> simultaneously.</p>
<p>I am trying to optimize  <code>Sample_th</code> and <code>Probe_th</code> with corresponding optimizer:</p>
<p><code>optimizer = torch.optim.Adam([ {'params': Probe_th, 'lr': 0.5e-1,'weight_decay':0.004}, {'params': Sample_th, 'lr':0.5e-1,'weight_decay':0.04}, ])</code></p>
<p>In the folllowing loop:</p>
<pre><code class="lang-auto">
for i in range(epoch_num):
    nums = list(range(len(measured_abs)))
    np.random.shuffle(nums)
    for num in nums:
        optimizer.zero_grad()
        Obj_th = Slices_th[num]
        Meas = Measured_th[num]
        loss = Sq_err(Propagation_model.forward(probe = Probe_th,obj = Obj_th,),Meas)
        err.append(float(loss.cpu()))
        loss.backward()
        optimizer.step()
        
    print(np.mean(err),'---',i)
    err_long.append(np.mean(err))
</code></pre>
<p>Currently, main probem is that during each  <code>optimizer.step()</code> gradient calculated for the whole <code>Sample_th</code> which takes most of the time, despite i know that only a small part of it corresponding to current slice participated in loss calculation and should be optimized during this step. From the other hand, I can’t  separate <code>Slices_th</code> into independent arrays, since I should take into account their mutual overlap, so change of one of them  during <code>optimizer.step()</code> should be spreaded among several of them correspondingly.</p>
<p>Sorry for such long explanation, I just don’t know how to explain it shorter)</p>