Multiprocessing using pytorch

username: Abueidda
cooked: <p>Hi,</p>
<p>I am trying to parallelize the following function:</p>
<pre><code class="lang-auto">def assemble(u):
    Fint=torch.zeros((ndof,1))
    Fext=torch.zeros((ndof,1))
    for iel in range(tne):
        elnodes=elems[iel,:].long() 
        xcel=nodes[elnodes,0] 
        ycel=nodes[elnodes,1]
        zcel=nodes[elnodes,2] 
        dof=torch.tensor([3*elnodes[0],3*elnodes[0]+1,3*elnodes[0]+2,\
                          3*elnodes[1],3*elnodes[1]+1,3*elnodes[1]+2,\
                          3*elnodes[2],3*elnodes[2]+1,3*elnodes[2]+2,\
                          3*elnodes[3],3*elnodes[3]+1,3*elnodes[3]+2,\
                          3*elnodes[4],3*elnodes[4]+1,3*elnodes[4]+2,\
                          3*elnodes[5],3*elnodes[5]+1,3*elnodes[5]+2,\
                          3*elnodes[6],3*elnodes[6]+1,3*elnodes[6]+2,\
                          3*elnodes[7],3*elnodes[7]+1,3*elnodes[7]+2]).flatten()
        u_el=u[dof]
        strain,stress = Stress_Strain(xcel,ycel,zcel,u_el)
        Fint_e,Fext_e = KEL(xcel,ycel,zcel,strain,stress)
        Fint[dof,0]  += Fint_e.squeeze(1)
        Fext[dof,0]  += Fext_e.squeeze(1)
    return Fint, Fext
</code></pre>
<p><strong>In particular, I am trying to parallelize the for loop in the function. Firstly, is it possible to do that? If yes, any advice/help would be much appreciated.</strong></p>

username: ptrblck
cooked: <p>Could you post a code snippet with random initializations for each tensor, please?<br>
This would make it easier to see, what shapes you are using and how we could potentially avoid the for loop.</p>

username: Abueidda
cooked: <p>Thanks for the reply, <a class="mention" href="/u/ptrblck">@ptrblck</a> . It would be great if the for loop can be avoided. The list of tensors used is below:</p>
<pre><code class="lang-auto">ndof    = 120000
tne     = 80000
u       = torch.rand(ndof)
elems   = torch.rand((tne, 8))
nodes   = torch.rand((ndof//3, 3))
elnodes = torch.rand(8)
xcel    = torch.rand(8)
ycel    = torch.rand(8)
zcel    = torch.rand(8)
dof     = torch.rand(24)
u_el    = torch.rand(24)
strain  = torch.rand((8, 6))
stress  = torch.rand((8, 6))
Fint_e  = torch.rand((24, 1))
Fext_e  = torch.rand((24, 1))
Fint    = torch.rand((ndof, 1))
Fext    = torch.rand((ndof, 1))
</code></pre>
<p>Thanks in advance</p>

username: ptrblck
cooked: <p>Thanks for the update.<br>
Would <code>Stress_Strain</code> and <code>KEL</code> accept batches or only the current input shapes?</p>

username: Abueidda
cooked: <p>Thanks again, <a class="mention" href="/u/ptrblck">@ptrblck</a>.<br>
The current forms of <code>Stress_Strain</code> and <code>KEL</code> do not accept batches. However, I feel that they can be changed, so they accept. Please see KEL below:</p>
<pre><code class="lang-auto">def KEL(xcel,ycel,zcel,strain,stress):
    eta_v = torch.tensor([-1/(3)**(0.5),1/(3)**(0.5)])
    psi_v = torch.tensor([-1/(3)**(0.5),1/(3)**(0.5)])
    phi_v = torch.tensor([-1/(3)**(0.5),1/(3)**(0.5)])
    intpt = torch.tensor([[psi_v[0],eta_v[0],phi_v[0]],
                          [psi_v[1],eta_v[0],phi_v[0]],
                          [psi_v[1],eta_v[1],phi_v[0]],
                          [psi_v[0],eta_v[1],phi_v[0]],
                          [psi_v[0],eta_v[0],phi_v[1]],
                          [psi_v[1],eta_v[0],phi_v[1]],
                          [psi_v[1],eta_v[1],phi_v[1]],
                          [psi_v[0],eta_v[1],phi_v[1]]])
    w=1 
    Fint_e=torch.zeros((ndofe,1))
    Fext_e=torch.zeros((ndofe,1))
    for i in range(8):
        psi,eta,phi=intpt[i,:]
        Bmat, Jdet=shpG(psi,eta,phi,xcel,ycel,zcel)
        W=w*Jdet
        Fint_e[:,0] += W*torch.mm(torch.t(Bmat), stress[i,:].unsqueeze(1)).squeeze(1)
    return Fint_e, Fext_e
</code></pre>
<p>Stress_Strain function has a similar fashion.</p>

username: Abueidda
cooked: <p>It seems that I found a way to make them accept batches. Thanks for the suggestion. <img src="https://discuss.pytorch.org/images/emoji/apple/grinning.png?v=9" title=":grinning:" class="emoji" alt=":grinning:"></p>