รู้จักกับ Generator ใน Python ของเทพ ๆ Optimise ได้มหาศาล

หลาย ๆ ครั้งเวลาเราเขียน Python โดยเราอาจจะเจอเคสที่เราจะต้องเขียน Function และมันจะต้องคืนค่ากลับมาเป็น Iterator ต่าง ๆ เป็น List หรือ Dict อะไรก็ตาม โดยอาจจะเอาไปใช้งานต่อใน Function ต่อไป ถ้าข้อมูลนั้นมีขนาดไม่ใหญ่มาก ก็อาจจะไม่มีปัญหาอะไรเลย เราก็แค่คืนค่ากลับไปที่ RAM ก็จบ แล้วเราก็เรียกต่อ ง่ายมาก ๆ เป็นเคสปกติเลย แต่ถ้า RAM เรามีน้อย หรือ Data ขนาดใหญ่ ๆ ละ Generator คือคำตอบของเรื่องนี้เลย

Generator คืออะไร ?

Generator เป็น Iterator ประเภทหนึ่ง สิ่งที่ทำมันคือ มันสามารถ Generate ค่าให้เราในระหว่างที่เราเรียกได้เลย กล่าวคือ เราไม่จำเป็นที่จะต้องรัน Function ให้จบเพื่อให้ได้ค่า แต่เราสามารถค่อย ๆ หยอดออกมาเมื่อมีคนเรียก Generator ขึ้นมาอาจจะผ่านการ Loop เข้าไป ซึ่งส่วนใหญ่แล้ว เรามักจะเจอ หรือสร้าง Generator ในรูปแบบของ Function ซะเยอะ

def square (num_list) :
	results = []
    for num in num_list :
    	results.append(num * num)
    return results

ด้านบนนี้เราลองสร้าง Function ง่าย ๆ ในการยกกำลังสองตัวเลขที่เราใส่เข้ามา แต่สังเกตุดูว่า เราใส่เลขมาเป็นชุด เป็น List ของ Integer สุดท้าย เราก็ให้มันหาคำตอบ เก็บลง List และคืน List กลับไป

>>> square([1,2,3,4])
[1, 4, 9, 16]
>>> type(square([1,2,3,4]))
<class 'list'>

ถ้าเราลองเรียก Function นี้ออกมาดู เราจะเห็นได้เลยว่า มันก็ตามนั้นนะ มันก็หาค่ายกำลังสองของตัวเลขแบบที่มันควรจะเป็น แบบที่เราต้องการ คืนกลับมาเป็น List ซึ่งเราอาจจะเอาไป Iterate เพื่อทำอะไรบางอย่างต่อได้ ในเคสนี้เราอาจจะไม่เห็นว่าทำไมมันเป็นปัญหา แต่ถ้าลองคิดดูว่า ถ้าเกิดตัวเลขนั้นมีสัก 10-100 ล้านตัว คิดว่ามันจะมหาศาลขนาดไหนกัน ทั้งเก็บเลข และ ไหนจะอ่านเข้ามา ไหนจะเก็บผลัพธ์อีก ก็คือ เครื่องแตกได้เลยนะ เพราะ RAM ไม่พอ เราลองมาเขียนให้มันเป็น Generator กันดีกว่า

def square_gen (num_list) :
	for num in num_list:
    	yield num * num

Code ด้านบนนี้ เราลองเขียนใหม่ ให้เป็น Generator ดู เราทำเหมือนเดิมเลยคือ เรารับค่าตัวเลขมา แล้ว เราก็ค่อย ๆ Loop ไล่ไปตามตัวเลขทีละตัว แต่สังเกตว่า เราจะไม่ได้ ทำทั้งหมดให้จบก่อน และ เราใส่คำสั่ง แปลก ๆ อย่าง yield ด้วย มันเป็นคำสั่งที่บอกให้ Python คืนค่ากลับไปเหมือน return ที่เราใช้ใน Function ทั่ว ๆ ไป แต่ yield เมื่อคืนค่ากลับไป มันจะรอเราเรียกเพื่อทำงานต่อ

>>> square_gen([1,2,3,4])
<generator object square_gen at 0x7f83881eadd0>

ถ้าเราลองเรียกเลยตรง ๆ เราจะเห็นว่า เราจะไม่ได้ค่าคำตอบคืนมาแบบที่เราเขียนเป็น Function ธรรมดา แต่เราจะได้กลับมาเป็น Generator Object แทนซะงั้น ซึ่งอย่างที่เราบอกว่า พวกนี้มันเป็น Iterator ดังนั้น เราสามารถ Loop เพื่อเข้าถึงค่าเหล่านี้ได้

square_list = square_gen([1,2,3,4])

for sq_num in square_list :
	print(sq_num)

ลองรันดู เราก็จะได้คำตอบเหมือนกับตอนที่เราเขียนเป็น Function ปกติทุกประการ แต่การมาได้ซึ่งคำตอบนี่แหละที่ต่างกัน ถ้าเราลองไปสังเกตเวลาเราเขียน Python จริง ๆ แล้ว เราก็ใช้ Generator กันเยอะมาก ๆ เลยเหมือนกัน ตัวอย่างเช่นเวลาเรา Loop เป็นตัวเลขลงมาเรื่อย ๆ เราจะชอบใช้ range() กันบ่อยมาก ๆ หรือแม้กระทั้ง enumerate() มันก็เป็น Generator เหมือนกัน ลองรันออกมาดู แล้วจะรู้เลย

มันต่างกันขนาดนั้นเลยเหรอ

ถึงแม้ว่าเราจะได้ผลลัพธ์เหมือนกัน แต่การมาได้ซึ่งผลลัพธ์ก็เป็นอีกสิ่งที่คำคัญไม่แพ้กัน เราลองมาดูกันดีกว่าว่า ถ้าเราเอาทั้ง 2 วิธีคือ การใช้ Function ปกติ กับการทำออกมาเป็น Generator มันจะต่างกันแค่ไหน ในการทดสอบนี้เราจะทดสอบเรื่องของเวลา โดยใช้ Standard Package time ใน Python และ การใช้ Memory ผ่าน Package memory_profiler กัน

import memory_profiler
import time

def square (num_list) :
    results = []
    for num in num_list :
        results.append(num * num)
    return results

bench_start = memory_profiler.memory_usage()
start_time = time.time()
    
square(range(10000000))
    
elapsed = time.time() - start_time
bench_end = memory_profiler.memory_usage()
mem_usage = bench_end[0] - bench_start[0]
    
print("Normal Function Took " + str(elapsed) + ' sec(s) and consumed ' + str(mem_usage) + ' Mb')

Code ด้านบนเป็นการทดสอบแรก เราใช้ Function เดิมเลยคือการหาค่าจากตัวเลขยกกำลังสอง ซึ่งอันนี้เราจะ ให้มันหาผลยกกำลังสอง ตั้งแต่ 0 - 9,999,999 เลย โดยวัดเวลาในการหา และ จำนวน RAM ที่ใช้ออกมา

Normal Function Took 0.5470151901245117 sec(s) and consumed 7.390625 MB

ผลที่ได้ อย่างที่เห็นเลย เราใช้เวลาอยู่ที่ครึ่งวินาที และใช้ RAM ที่ประมาณ 7 MB เราลองไปดูฝั่ง Generator กันบ้างดีกว่า

import memory_profiler
import time

def square_gen (num_list) :
    for num in num_list:
    	yield num * num

bench_start = memory_profiler.memory_usage()
start_time = time.time()
    
square_gen(range(10000000))
    
elapsed = time.time() - start_time
bench_end = memory_profiler.memory_usage()
mem_usage = bench_end[0] - bench_start[0]
    
print("Generator Took " + str(elapsed) + ' sec(s) and consumed ' + str(mem_usage) + ' Mb')

ส่วนของ Generator เราทำเหมือนเดิมเลย แค่เปลี่ยน Function ให้เป็น Generator แทน ทดสอบเหมือนกันทุกประการ

Generator Took 3.0994415283203125e-05 sec(s) and consumed 0.03125 Mb

พอมาเป็น Generator เราจะเห็นได้เลยว่า เวลามันลดลงอย่างมหาศาลมาก ๆ เหลือเสี้ยววินาทีเท่านั้น กับใช้ RAM ไม่ถึงครึ่ง MB เลย แต่อย่าลืมนะว่า การที่เราใช้เป็น Generator และเราเรียกแค่นี้ มันยังไม่ได้รันเอาข้อมูลที่ต้องคำนวณออกมาเลย เราลองทำอะไรเพิ่มดูด้วยการเพิ่มส่วนของการหาผลบวกของ เลขยกกำลังสองทั้งหมดกัน

sum = 0
for result in results :
	sum += result

สิ่งที่เราลองเพิ่มคือ เราลองส่วนที่รับค่าแล้วเอามาบวกกับตัวแปร sum เพื่อให้เราได้ผลบวกของค่ายกกำลังสองทั้งหมด ลองดูผลกัน

Normal Function Took 1.1276490688323975 sec(s) and consumed 25.03125 MB
Generator Took 1.067375898361206 sec(s) and consumed 17.5625 MB

จากผล เราจะเห็นได้เลยว่า เวลาในการทำงานถือว่าใกล้เคียงกันมาก ๆ ห่างกันน้อยมาก ๆ แต่ดูจำนวน Memory ที่ใช้ เราได้ผลลัพธ์เหมือนกัน แต่การใช้ Memory ค่อนข้างต่างอย่างเห็นได้ชัดเลย โดยที่ตัว Generator ใช้ Memory น้อยกว่า เป็นเพราะ เราไม่ต้องเก็บผลจำนวน 10 ล้านไว้ก่อน เราค่อย Generate ระหว่างที่เรา Loop เลย ทำให้ใช้ Memory น้อยกว่านั่นเอง

ลองประโยชน์ที่ชัดเจนขึ้น

การที่เรายกตัวอย่างอันนี้อาจจะเห็นภาพได้ไม่ชัดเท่าไหร่เพราะขนาดของข้อมูลไม่ได้ใหญ่มาก ๆ เราขอยกตัวอย่าง Generator เวลาเราทำงานจริง ๆ ที่เห็นผลจริง ๆ สัก 2 อย่างดีกว่า อย่างแรก ถ้าเราเขียนพวก Tensorflow มาก่อน และทำงานกับข้อมูลขนาดใหญ่ เช่น 100 GB หรือมากกว่านั้น และ RAM ในเครื่องเรามีไม่ถึง การที่เราจะโหลดข้อมูลทั้งหมดขึ้นมาและอัดให้ Tensorflow ไป Pre-Process และ Train เลยอาจจะเป็นเรื่องที่เป็นไปได้ยาก (Paging ก็ได้อยู่แหละ แต่โคตรช้าเลย) ดังนั้นการอัดผ่าน Generator เพื่อให้มันจัดการ Feature ของเราระหว่างตอนที่เรา Train เลย มันก็ทำให้เราประหยัด RAM ขึ้นได้เยอะ (ลองดูใน Document ของ Tensorflow ได้) จากเดิมที่การโหลดข้อมูลขนาดใหญ่ทั้งหมดเลย มันเป็นเรื่องที่เป็นไปไม่ได้เลย แต่การใช้ Generator ทำให้การ Train ข้อมูลขนาดใหญ่ในเครื่องคอมพิวเตอร์ที่มี RAM น้อยเป็นไปได้

หรือในงานตอนที่ทำ Thesis เราทำงานกับ Whole Genome Sequencing ของคน ซึ่งขนาดของ Dataset ต่อ Sample ใหญ่มาก ๆ อาจจะใหญ่ถึงตัวอย่างละ 100 GB ++ เลยก็ได้ ถ้าเรามีสัก 10 ตัวอย่างโหลดเข้าไป นึกสภาพว่า ถ้าไม่มี Generator มันจะได้ขนาดไหนกัน

สรุป

Generator เป็น Iterator ที่ทำให้เราสามารถค่อย ๆ Generate ค่าออกมาได้เรื่อย ๆ เมื่อเราเรียกมันออกมา ทำให้เราสามารถ Optimise การใช้ Memory ได้ จากในการทดลองในบทความนี้ ทำให้เราเห็นเลยว่า การใช้งาน Generator ช่วยทำให้เราประหยัด Memory ในการรันได้เยอะพอตัวเลย ยิ่งถ้าเราทำงานกับข้อมูลขนาดใหญ่ ๆ ด้วยละก็ การใช้ Generator ยิ่งเป็นผลดีต่อ Performance มาก ๆ ทำให้เรามีพื้นที่ RAM เหลือไปเก็บอย่างอื่นได้อีกเยอะ