大栅格数据如何更快运算
背景介绍
这两周我在使用python进行大量的栅格数据的运算，在运算过程中遇到了数据量超级大但算力不足的问题。通过这两周的探索，也慢慢找到了一些加快栅格数据计算的方法，和读者分享。
原理
首先说一下加快栅格数据计算的基本准则：
（1）尽可能榨干电脑的全部性能，把电脑CPU、内存、固态硬盘、机械硬盘进行合理分工等；
（2）使用多线程运算；
（3）及时释放计算中占用的内存空间。
问题与解决方法
（1）数据量超过电脑内存，使用分块运算
在计算栅格数据时，是把数据放到内存中进行计算，如果栅格计算数量巨大，会爆内存。
分块方法就是采用横纵切割原始栅格，最后再将数据拼接起来。比如下面这个代码，通过RasterXSize和RasterYSize获取数据的大小,然后将栅格影像分为了4乘4，共计16块小栅格。
src_ds = gdal.Open(TIF_PATH)
x_size = src_ds.RasterXSize
y_size = src_ds.RasterYSize
# 修改分块大小，将图像分为更多的网格，以减小每个进程的内存需求
x_block_size = x_size // 4
y_block_size = y_size // 4
for x_offset in tqdm(range(0, x_size, x_block_size), desc=f"Processing {subfolder} x_block_size"):
for y_offset in range(0, y_size, y_block_size):
block_width = min(x_block_size, x_size - x_offset)
block_height = min(y_block_size, y_size - y_offset)
通过分块，能有效降低栅格运算过程中的内存占用情况。
（2）分块运算还是超过内存，使用mmap_array数组的运算
如果分块运行还是超过内存大小，这个时候就需要考虑将分块数据的中间数据存在硬盘中，等需要的时候再去读取。栅格的运算一般使用的是numpy模块，然后将数据转为数组array放到内存中计算。但如果你的栅格数据过大，就需要用到mmap_array，这是一个内存映射数组，可以保存到硬盘中。
array：array 是一个普通的 NumPy 数组，它是 numpy.ndarray 类的一个实例。这种数组将其数据直接存储在内存中。普通的 NumPy 数组用于处理可以容纳在内存中的数据集，并且在大多数情况下，计算和操作速度更快。然而，它们不能用于处理比可用内存更大的数据集。
mmap_array：mmap_array 是一个内存映射文件 (memory-mapped file) 数组，它是 numpy.memmap 类的一个实例。这种数组的数据存储在磁盘上的一个文件中，而不是直接存储在内存中。numpy.memmap 的主要优点是，它允许您处理比可用内存更大的数据集，因为数据只在需要时才从磁盘加载到内存中。对于非常大的数组或在多进程环境下共享数据时，这种方法非常有用。
因为mmap_array只是中间数据，应该记得运行一次，清理一次，防止占用内存。
#清理mmap_array
mmap_array.flush()
mmap_array._mmap.close()
（3）数据读取写入速度慢，在固态硬盘中运行
刚开始我使用了机械硬盘作为数据的运行盘和数据保存盘。
但硬盘的写入速度和读取速度经常爆100%，这个时候就知道了mmap_array数组需要和内存进行快速的读取和写入，由于mmap_array数组默认是保存到python脚本的同级目录之下，所以为了突破硬盘的限制，将脚本转移到固态硬盘中运行。
但还需要注意个问题，如果你运行的是超级大的栅格数据，固态硬盘的容量应该是不够保存的，因此每次运行完栅格数据后，都应该及时转移数据到机械硬盘。转移数据可以使用shutil模块，比如：
import shutil
#移动文件到机械硬盘
destination_file = os.path.join(destination_folder, f"{subfolder}.tif")
shutil.move(output_file, destination_file)
（4）数据量太大，使用数据压缩技术
这里要提及一下，最开始我使用过arcgis pro自带的arcpy进行数据计算，但arcpy数据生成结果是没有被压缩过，每一期的数据都会生成200G大小的栅格数据。
但是转为使用gdal模块后，输出数据的详细参数我可以直接控制，因此将输出的栅格数据进行DEFLATE压缩。为什么选择DEFLATE压缩？我这里考虑的是使用无损压缩、压缩率较高。
1682354050959.png
LZW压缩也可以满足大部分人的需求，可以根据自己的需求选择压缩。在选择压缩后，输出的数据需要注意，如果栅格数据大于4g，需要额外设置一个bigtif参数。下面是压缩代码：
TIF_PATH = driver.Create(output_file, x_size, y_size, 1, gdal.GDT_Int16, options=["COMPRESS=DEFLATE", "BIGTIFF=YES"])
（5）电脑性能无法发挥，善用多线程榨干电脑性能
在刚开始运行数据时，遇到一个问题：CPU无法发挥其全部实力，但内存已经满了。
比如我在运行过程中，就遇到CPU只占用了10%出头，但内存已经爆了。
有没有办法既提高CPU的运行速度，也不爆内存，还能提高运算速度？可以，使用多线程。
Python的多线程技术可以使用内置的 threading 模块来实现。比如：
import threading
def worker(num):
"""线程工作的函数"""
print(f"Thread-{num} is running.")
# 在这里添加工作内容 ...
print(f"Thread-{num} is finished.")
# 创建两个线程并启动
thread1 = threading.Thread(target=worker, args=(1,))
thread1.start()
thread2 = threading.Thread(target=worker, args=(2,))
thread2.start()
我们定义了一个 worker() 函数作为线程的工作内容。然后使用 threading.Thread() 创建两个线程，并分别传入工作函数 worker() 和参数 args=(1,) 或 args=(2,)。接着调用 start() 方法启动线程。
使用多线程，但如何才不能爆内存了？可以通过调整分块的大小，分块越小，内存占用越小，能带动的的线程数量越多。
但是分块的大小不是越小越好，会有一个阈值。这个需要根据电脑的性能来具体操作，总体而言，就是考虑到CPU占用、内存占用情况、分块大小进行动态调整。我画个简单的示意图：
代码示例
在这个代码中，我使用了分块技术进行栅格的运算，使用mmap_array存储中间数据映射内存文件，同时考虑到固态硬盘容量有限进行了数据转移，也使用了多线程技术达到了电脑的性能瓶颈。
该计算多期数据量超大的栅格平均值代码，这个代码不仅能处理栅格预算，也可以进行裁剪、重分类、镶嵌等，只需要把里面的功能换一换，自己调整一下参数便可以用来处理数据量超大的栅格数据。
import os
import numpy as np
from osgeo import gdal
from tqdm import tqdm
import gc
from multiprocessing import Pool
import uuid
import shutil
root_folder = r"E:\\cloud\\input\\"
output_folder = r"D:\\MIDDATA\\"
# 设置移动目标文件夹路径
destination_folder = r"E:\\cloud\\output\\"
# root_folder = r"D:\\keshan\\gdal\\代求\\"
# output_folder = r"D:\\keshan\\gdal\\结果\\"
def process_subfolder(subfolder):
try:
subfolder_path = os.path.join(root_folder, subfolder)
if os.path.isdir(subfolder_path):
files = [f for f in os.listdir(subfolder_path) if f.endswith('.tif')]
first = os.path.join(subfolder_path, files[0])
src_ds = gdal.Open(first.encode('utf-8'))
x_size = src_ds.RasterXSize
y_size = src_ds.RasterYSize
data_type = src_ds.GetRasterBand(1).DataType
projection = src_ds.GetProjection()
geotransform = src_ds.GetGeoTransform()
output_file = os.path.join(output_folder, f"{subfolder}.tif")
driver = gdal.GetDriverByName("GTiff")
dst_ds = driver.Create(output_file, x_size, y_size, 1, gdal.GDT_Int16, options=["COMPRESS=DEFLATE", "BIGTIFF=YES"])
dst_ds.SetProjection(projection)
dst_ds.SetGeoTransform(geotransform)
# 修改分块大小，将图像分为更多的网格，以减小每个进程的内存需求
x_block_size = x_size // 29
y_block_size = y_size // 29
for x_offset in tqdm(range(0, x_size, x_block_size), desc=f"Processing {subfolder} x_block_size"):
for y_offset in range(0, y_size, y_block_size):
block_width = min(x_block_size, x_size - x_offset)
block_height = min(y_block_size, y_size - y_offset)
# dtype = np.int32
# mmap_array = np.memmap('mmap.bin', dtype=dtype, mode='w+', shape=(block_height, block_width))
dtype = np.int32
# 为每个进程生成一个唯一的内存映射文件名
mmap_filename = f'mmap_{uuid.uuid4().hex}.bin'
mmap_array = np.memmap(mmap_filename, dtype=dtype, mode='w+', shape=(block_height, block_width))
for file_name in files:
file_path = os.path.join(subfolder_path, file_name)
src_ds = gdal.Open(file_path)
array = src_ds.GetRasterBand(1).ReadAsArray(x_offset, y_offset, block_width, block_height)
array = array.astype('int32')
mmap_array += array
src_ds = None
mean_array = (mmap_array / len(files)).astype(np.int16)
out_band = dst_ds.GetRasterBand(1)
out_band.WriteArray(mean_array, xoff=x_offset, yoff=y_offset)
mmap_array.flush()
mmap_array._mmap.close()
os.unlink(mmap_filename)
del mmap_array, mean_array, out_band, array
gc.collect()
dst_ds = None
#移动文件到机械硬盘
destination_file = os.path.join(destination_folder, f"{subfolder}.tif")
shutil.move(output_file, destination_file)
#保存信息
error_message = f"{output_file} GO_sucess: "
print(error_message)
# 将成功信息保存到log.txt
error_log_path = r"E:\cloud\output\log.txt"
with open(error_log_path, "a") as error_log:
error_log.write(error_message + "\n")
except Exception as e:
error_message = f"{subfolder} GO_WRONG: {str(e)}"
print(error_message)
# 将错误信息保存到log.txt
error_log_path = os.path.join(destination_folder, "log.txt")
with open(error_log_path, "a") as error_log:
error_log.write(error_message + "\n")
if __name__ == "__main__":
subfolders = [subfolder for subfolder in os.listdir(root_folder) if os.path.isdir(os.path.join(root_folder, subfolder))]
# 设置进程池的大小为28，即32线程的93%
num_processes = 10
with Pool(num_processes) as pool:
# 使用tqdm显示进度
list(tqdm(pool.imap(process_subfolder, subfolders), total=len(subfolders), desc="Processing subfolders"))
总结
之前我没有处理过这么庞大的数据，以为会很简单，但等真正投身进去，才知道提高栅格的数据也是需要技巧的。下面是我的技巧总结：
（1）使用多线程提高计算速度；
（2）测试自己的电脑瓶颈在哪，比如内存、还是硬盘速度、cpu速度？cpu有空余，增加线程数量；内存有空余，数据切片大一些。
（3）固态硬盘用来存放中间文件mmap，固态硬盘不够大，可以像我一样，把生成文件移动到机械硬盘中去
（4）tif文件超过4G，要记得gdal导出栅格时参数设置为“BIGTIFF=YES”
（5）栅格分块跑数据。
（6）输出数据尽可能使用压缩功能，这里建议使用DEFLATE压缩。
（7）栅格的像素深度能用int类型就不用float，能用int16就不用int64。
（8）尽量使用gdal库跑数据，其他的地理空间第三方库比如arcpy有很多参数无法自己调整，比如指定导出的栅格压缩类型。
腾讯云开发者
