microsoft / speedyrec Goto Github PK

    pool = Pool(processes=args.world_size)
    results = []
    sigle_size = news_num//args.world_size
    for rank in range(args.world_size):
        start = sigle_size*rank
        end = sigle_size*(rank+1)
        if rank == args.world_size -1:
            end = news_num
        local_features = news_feature[start:end]
        result = pool.apply_async(sigle_process_infer, args=(rank, local_features, checkpoint, args))
        results.append(result)
    pool.close()
    pool.join()

    results = [x.get() for x in results]
    news_vecs = np.concatenate(results,0)