使用MPI並行化遺傳演算法框架GAFT

PytLab醬發表於2019-02-12

前言

本文中作者使用MPI的Python介面mpi4py來將自己的遺傳演算法框架GAFT進行多程式並行加速。並對加速效果進行了簡單測試。

專案連結:

GitHub: github.com/PytLab/gaft
PyPI: pypi.python.org/pypi/gaft

正文

我們在用遺傳演算法優化目標函式的時候，函式通常都是高維函式，其導數一般比較難求取。這樣我們的適應度函式計算通常都是比較費時的計算。

例如在使用遺傳演算法尋找最優結構時候通常需要呼叫量化軟體進行第一性原理計算結構的total energy，這是非常費時的過程; 例如我們優化力場引數的時候，以力場計算出的能量同基準能量之前的誤差作為適應度，也需要呼叫相應的力場程式獲取總能量來求取，同樣這個過程也是相對耗時的。

這就會導致一個問題，當我們的種群比較大的時候，我們需要利用適應度資訊來產生下一代種群，這時候每一代繁殖的過程將會很耗時。但有幸的是，種群的選擇交叉變異過程對於種群中的個體都是相互獨立的過程，我們可以將這一部分進行並行處理來加速遺傳演算法的迭代。

使用mpi4py

由於實驗室的叢集都是MPI環境，我還是選擇使用MPI介面來將程式碼並行化，這裡我還是用了MPI介面的Python版本mpi4py來將程式碼並行化。關於mpi4py的使用，我之前寫過一篇部落格專門做了介紹，可以參見《Python多程式並行程式設計實踐-mpi4py的使用》

將mpi4py的介面進一步封裝

為了能讓mpi的介面在GAFT中更方便的呼叫，我決定將mpi4py針對遺傳演算法中需要用的地方進行進一步封裝，為此我單獨寫了個MPIUtil類, 詳細程式碼參見gaft/mpiutil.py。

封裝通訊子常用的介面

例如程式同步, 獲取rank，程式數，判斷是否為主程式等。

class MPIUtil(object):
    def __init__(self):
        logger_name = 'gaft.{}'.format(self.__class__.__name__)
        self._logger = logging.getLogger(logger_name)

    # Wrapper for common MPI interfaces.
    def barrier(self):
        if MPI_INSTALLED:
            mpi_comm = MPI.COMM_WORLD
            mpi_comm.barrier()

    @property
    def rank(self):
        if MPI_INSTALLED:
            mpi_comm = MPI.COMM_WORLD
            return mpi_comm.Get_rank()
        else:
            return 0

    @property
    def size(self):
        if MPI_INSTALLED:
            mpi_comm = MPI.COMM_WORLD
            return mpi_comm.Get_size()
        else:
            return 1

    @property
    def is_master(self):
        return self.rank == 0複製程式碼

組內集合通訊介面

由於本次並行化的任務是在種群繁衍時候進行的，因此我需要將上一代種群進行劃分，劃分成多個子部分，然後在每個程式中對劃分好的子部分進行選擇交叉變異等遺傳操作。在最後將每個字部分得到的子種群進行收集合並。為此寫了幾個劃分和收集的介面:

    def split_seq(self, sequence):
        '''
        Split the sequence according to rank and processor number.
        '''
        starts = [i for i in range(0, len(sequence), len(sequence)//self.size)]
        ends = starts[1: ] + [len(sequence)]
        start, end = list(zip(starts, ends))[self.rank]

        return sequence[start: end]

    def split_size(self, size):
        '''
        Split a size number(int) to sub-size number.
        '''
        if size < self.size:
            warn_msg = ('Splitting size({}) is smaller than process ' +
                        'number({}), more processor would be ' +
                        'superflous').format(size, self.size)
            self._logger.warning(warn_msg)
            splited_sizes = [1]*size + [0]*(self.size - size)
        elif size % self.size != 0:
            residual = size % self.size
            splited_sizes = [size // self.size]*self.size
            for i in range(residual):
                splited_sizes[i] += 1
        else:
            splited_sizes = [size // self.size]*self.size

        return splited_sizes[self.rank]

    def merge_seq(self, seq):
        '''
        Gather data in sub-process to root process.
        '''
        if self.size == 1:
            return seq

        mpi_comm = MPI.COMM_WORLD
        merged_seq= mpi_comm.allgather(seq)
        return list(chain(*merged_seq))複製程式碼

用於限制程式在主程式執行的裝飾器

有些函式例如日誌輸出，資料收集的函式，我只希望在主程式執行，為了方便，寫了個裝飾器來限制函式在主程式中執行:

def master_only(func):
    '''
    Decorator to limit a function to be called
    only in master process in MPI env.
    '''
    @wraps(func)
    def _call_in_master_proc(*args, **kwargs):
        if mpi.is_master:
            return func(*args, **kwargs)

    return _call_in_master_proc複製程式碼

在遺傳演算法主迴圈中新增並行

主要在種群繁衍中對種群針對程式數進行劃分然後並行進行遺傳操作併合並子種群完成並行，程式碼改動很少。詳見:github.com/PytLab/gaft…

# Enter evolution iteration.
for g in range(ng):
    # Scatter jobs to all processes.
    local_indvs = []
    local_size = mpi.split_size(self.population.size // 2)

    # Fill the new population.
    for _ in range(local_size):
        # Select father and mother.
        parents = self.selection.select(self.population, fitness=self.fitness)
        # Crossover.
        children = self.crossover.cross(*parents)
        # Mutation.
        children = [self.mutation.mutate(child) for child in children]
        # Collect children.
        local_indvs.extend(children)

    # Gather individuals from all processes.
    indvs = mpi.merge_seq(local_indvs)
    # The next generation.
    self.population.individuals = indvs複製程式碼

測試加速效果

測試一維搜尋

下面我針對專案中的一維優化的例子進行並行加速測試來看看加速的效果。例子程式碼在/examples/ex01/

由於自己本子核心數量有限，我把gaft安裝在實驗室叢集上使用MPI利用多核心進行平行計算以為優化，種群大小為50，代數為100，針對不同核心數可以得到不同的優化時間和加速比。視覺化如下圖:

核心數	優化時間(s)	加速比
1	1.473	1.0
2	0.877	1.68
3	0.657	2.24
4	0.533	2.76
5	0.467	3.15
6	0.540	2.73
7	0.431	3.42
8	0.382	3.86
9	0.355	4.15
10	0.317	4.65

核心數與優化時間的關係:

核心數與加速比:

測試力場優化

這裡我對自己要研究的物件進行加速測試，這部分程式碼並未開源，針對每個個體的適應度計算都需要呼叫其他的計算程式，因此此過程相比直接有函式表示式的目標函式計算要耗時很多。

同樣，我針對不同核心數看看使用MPI在叢集上加速的效果:

核心數	優化時間(s)	優化時間	加速比
1	2.29e04	6 h 21 min	1.0
2	1.94e04	5 h 23 min	1.18
4	1.62e04	4 h 30 min	1.41
6	1.35e04	3 h 45 min	1.69
12	8.73e03	2 h 25 min	2.62