Multiprocessing using pytorch

1

Hi,

I am trying to parallelize the following function:

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

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.

2

Could you post a code snippet with random initializations for each tensor, please?
This would make it easier to see, what shapes you are using and how we could potentially avoid the for loop.

1 Like
3

Thanks for the reply, @ptrblck . It would be great if the for loop can be avoided. The list of tensors used is below:

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))

Thanks in advance

4

Thanks for the update.
Would Stress_Strain and KEL accept batches or only the current input shapes?

5

Thanks again, @ptrblck.
The current forms of Stress_Strain and KEL do not accept batches. However, I feel that they can be changed, so they accept. Please see KEL below:

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

Stress_Strain function has a similar fashion.

6

It seems that I found a way to make them accept batches. Thanks for the suggestion. :grinning: