使用Pytorch加速PSO算法？

Z时代
2024-01-10
分类：问答

我已经实施了PSO算法使用PyTorch,torch.cuda.FloatTensor。我设法加速了2次，但我期待的不止于此。它认为步骤检查限制(line 41-55)和参数更新(line 58-72)是原因。有一个“GPU任何”这样做？使用Pytorch加速PSO算法？

有是我的代码：

import torch 
import numpy as np 
import matplotlib.pyplot as plt 
dtype = torch.cuda.FloatTensor 
def fitness(x,y): 
    return -torch.abs(torch.sin(x)*torch.cos(y)*torch.exp(torch.abs(1-(torch.sqrt(x**2+y**2))/(3.1415)))) 
def velocity(v, gxbest, pxbest, pybest, x, pop, w, c1, c2): 
    return w*torch.rand(pop).type(dtype)*v + \ 
     c1*torch.rand(pop).type(dtype)*(pxbest - x) + \ 
     c2*torch.rand(pop).type(dtype)*(gxbest.expand(x.size(0)) - x) 
def PSO(pop, xmax, xmin, niter, wmax, wmin, c1, c2): 
    vx     = torch.rand(pop).type(dtype) 
    vy     = torch.rand(pop).type(dtype) 
    best    = np.zeros(niter) 
    x     = (xmax - xmin)*torch.rand(pop).type(dtype) + xmin 
    y     = (xmax - xmin)*torch.rand(pop).type(dtype) + xmin 
    z     = fitness(x,y) 
    [minz, indexminz] = z.min(0) 
    gxbest   = x[indexminz] 
    gybest   = y[indexminz] 
    pxbest   = x 
    pybest   = y 
    pzbest   = z 
    for K in range(niter): 
     w  = wmax - ((wmax - wmin)/niter) * (K) 
     vnextx = velocity(vx, gxbest, pxbest, pybest, x, pop, w, c1, c2) 
     xnext = x + vnextx 
     vnexty = velocity(vy, gxbest, pxbest, pybest, y, pop, w, c1, c2) 
     ynext = y + vnexty 
     **(41)** xnext = xnext.cpu() 
     ynext = ynext.cpu() 
     idxmax  = (xnext > xmax) # elements that are bigger that upper limit 
     idxmim  = (xnext < xmin) # elements that are smaller that upper limit 
     xnext[idxmax] = xmax 
     xnext[idxmim] = xmin 
     idymax  = (ynext > xmax) # elements that are bigger that upper limit 
     idymim  = (ynext < xmin) # elements that are smaller that upper limit 
     ynext[idymax] = xmax 
     ynext[idymim] = xmin 
     xnext = xnext.cuda() 
     **(55)** ynext = ynext.cuda() 
     znext = fitness(xnext,ynext) 
     **(58)**[minznext, indexminznext] = znext.min(0) 
     if (minznext[0] < minz[0]): 
      minz = minznext 
      gxbest = xnext[indexminznext] 
      gybest = ynext[indexminznext] 
     indexpbest   = (znext < pzbest) 
     pxbest[indexpbest] = xnext[indexpbest] 
     pybest[indexpbest] = ynext[indexpbest] 
     pzbest[indexpbest] = znext[indexpbest] 
     x     = xnext 
     y     = ynext 
     vx     = vnextx 
     **(72)** vy     = vnexty 
     best[K] = minz.cpu().numpy() 
    return gxbest, gybest , minz, best 
def main(): 
    pop, xmax, xmin, niter = 10000, 10, -10, 10 
    wmax = 0.9 
    wmin = 0.4 
    c1 = 2.05 
    c2 = 2.05 
    xbest, ybest, fitbest, best = PSO(pop, xmax, xmin, niter, wmax, wmin, c1, c2) 
    print(xbest) 
    print(ybest) 
    print(fitbest) 
    t = np.linspace(0,niter,niter) 
    plt.plot(t, best, 'k.-') 
    plt.show() 
main()