ใช้ Google Colab TPU สำหรับงาน Deep learning เพื่อการทำงานที่เร็วขึ้น

4 min readAug 25, 2020

ที่ผ่านมาเป็นความคุ้นเคยกันดีอยู่แล้ว สำหรับการ Training deep learning model ทั้งบน CPU และ GPU ซึ่งแบบหลังมีความนิยมมากกว่า เพราะเนื่องจากความเร็วในการประมวลผลเชิงคณิตศาสตร์ของ GPU นำหน้า CPU ไปเยอะมาก แต่ด้วยความที่ AI หรือ Neural network ก็พัฒนาความซับซ้อนเพิ่มขึ้นไปอีก ดังนั้นการประมวลผลแบบเดิม จึงอาจให้ความเร็วที่ไม่น่าพอใจมากนัก

อาจเคยได้ยินมาสักพักแล้ว สำหรับการประมวลแบบในรูปแบบของ TPU (Tensor Processing Unit) ซึ่งเป็น Open source จากทาง Google เพื่อการประมวลสำหรับ Machine learning โดยเฉพาะ ซึ่งเทคโนโลยีถูกออกแบบมาให้ใช้พลังงานในการประมวลต่ำกว่า CPU และ GPU แต่ไม่ได้นำมาทดแทนการทำงานของทั้งคู่แต่อย่างใด

CPU vs GPU vs TPU

ก่อนที่วันนี้จะมาลองใช้งาน TPU ขอทำความรู้จักกับการประมวลผลทั้ง 3 รูปแบบกันก่อน โดยไม่ขอลงรายละเอียดเกี่ยวกับ Hardware architecture ขอเริ่มจาก CPU (Central Processing Unit) เพราะเป็นสิ่งที่ถูกรู้จักกันมากที่สุด เพราะคือหน่วยประมวลผลการของคอมพิวเตอร์ ที่ใช้ควบคุมสิ่งต่าง ๆ ที่เกิดขึ้นภายในคอมพิวเตอร์ทั้ง Input และ Output รวบไปถึงการคำนวณเชิงคณิตศาสตร์พื้นฐาน จึงทำให้ CPU สามารถมาทำใช้ในงาน Machine learning ได้บ้าง

ต่อมาคือ GPU (Graphic Processing Unit) พิจารณาจากชื่อแล้ว มันเป็นหน่วยประมวลผลเกี่ยวกับกราฟฟิกโดยเฉพาะ ซึ่งเบื้องหลังกราฟฟิกที่เกิดขึ้นในคอมพิวเตอร์ ล้วนแล้วแต่เป็นคณิตศาสตร์ ดังนั้นจำเป็นคำตอบที่ว่าทำไม GPU จึงสามารถนำมาใช้งานการ Machine learning ได้อย่างรวดเร็วและมีประสิทธิภาพกว่า CPU

สุดท้าย TPU อย่างที่บอกไปว่าเป็นเครื่องที่ Google พัฒนาขึ้นมาเอง เพื่อให้เหมาะกับการทำงานของ Deep learning library อย่าง Tensorflow เมื่อ TPU ถูกออกแบบมาให้ทำงานเฉพาะทาง หมายความว่าระบบต้องมีความซับซ้อนน้อยลง (ไม่ต้องประมวลผลอย่างอื่น คิดเลขเพียงอย่างเดียว) ส่งผลให้สามารถทำงานได้รวดเร็วยิ่งขึ้น ซึ่งให้งานด้าน Machine learning มีประสิทธิภาพมากยิ่งขึ้นนั่นเอง

Using TPU on Google Colab

มาถึงการทดลองใช้ TPU สำหรับงาน Deep learning ซึ่ง TPU Option สามารถเลือกเปิดใช้งานได้จาก Google Colab ที่เมนู Runtime > Change runtime type ตรง Hardware accelerator ให้เลือกเป็น TPU

ขอใช้ Convolutional neural network ที่เคยทำไว้ตอนทำ Classification model ของฟิล์ม Chest X-Ray COVID-19 ซึ่งตอนนั้นได้ใช้ GPU ในการ Training model เพื่อให้เห็นความแตกต่าง จึงใช้ Parameters ทุกอย่างเหมือนเดิม

Chest X-Ray COVID-19 Model ตอนที่ 2: สร้าง CNN Model from scratch

ได้เวลา Training CNN Model แบบตั้งแต่ 0 กันแล้ว

medium.com

Epoch 1/20
81/81 [==============================] - 61s 750ms/step - loss: 2.0039 - accuracy: 0.9023 - val_loss: 0.3126 - val_accuracy: 0.9152
Epoch 2/20
81/81 [==============================] - 57s 698ms/step - loss: 0.1075 - accuracy: 0.9651 - val_loss: 0.1252 - val_accuracy: 0.9621
Epoch 3/20
81/81 [==============================] - 56s 697ms/step - loss: 0.0599 - accuracy: 0.9806 - val_loss: 0.1016 - val_accuracy: 0.9693
Epoch 4/20
81/81 [==============================] - 56s 696ms/step - loss: 0.0384 - accuracy: 0.9876 - val_loss: 0.1186 - val_accuracy: 0.9639
Epoch 5/20
81/81 [==============================] - 56s 697ms/step - loss: 0.0140 - accuracy: 0.9953 - val_loss: 0.1009 - val_accuracy: 0.9639
Epoch 6/20
81/81 [==============================] - 56s 697ms/step - loss: 0.0197 - accuracy: 0.9953 - val_loss: 0.1009 - val_accuracy: 0.9711
Epoch 7/20
81/81 [==============================] - 56s 697ms/step - loss: 0.2435 - accuracy: 0.9744 - val_loss: 0.4511 - val_accuracy: 0.9729
Epoch 8/20
81/81 [==============================] - 56s 686ms/step - loss: 4.5123 - accuracy: 0.9419 - val_loss: 0.5956 - val_accuracy: 0.9404
Epoch 9/20
81/81 [==============================] - 55s 682ms/step - loss: 0.1296 - accuracy: 0.9752 - val_loss: 0.2419 - val_accuracy: 0.9549
Epoch 10/20
81/81 [==============================] - 56s 687ms/step - loss: 0.0560 - accuracy: 0.9891 - val_loss: 0.2152 - val_accuracy: 0.9603
Epoch 11/20
81/81 [==============================] - 56s 685ms/step - loss: 0.0337 - accuracy: 0.9915 - val_loss: 0.2176 - val_accuracy: 0.9657
Epoch 12/20
81/81 [==============================] - 56s 686ms/step - loss: 0.0856 - accuracy: 0.9822 - val_loss: 0.3647 - val_accuracy: 0.9585
Epoch 13/20
81/81 [==============================] - 55s 684ms/step - loss: 0.0526 - accuracy: 0.9868 - val_loss: 0.2556 - val_accuracy: 0.9693
Epoch 14/20
81/81 [==============================] - 55s 685ms/step - loss: 0.0283 - accuracy: 0.9953 - val_loss: 0.2608 - val_accuracy: 0.9675
Epoch 15/20
81/81 [==============================] - 55s 683ms/step - loss: 0.0142 - accuracy: 0.9969 - val_loss: 0.3273 - val_accuracy: 0.9603
Epoch 16/20
81/81 [==============================] - 56s 685ms/step - loss: 0.0179 - accuracy: 0.9938 - val_loss: 0.2008 - val_accuracy: 0.9693
Epoch 17/20
81/81 [==============================] - 55s 684ms/step - loss: 0.0226 - accuracy: 0.9938 - val_loss: 0.2208 - val_accuracy: 0.9711
Epoch 18/20
81/81 [==============================] - 56s 686ms/step - loss: 0.0105 - accuracy: 0.9977 - val_loss: 0.2050 - val_accuracy: 0.9657
Epoch 19/20
81/81 [==============================] - 55s 685ms/step - loss: 0.0124 - accuracy: 0.9977 - val_loss: 0.2784 - val_accuracy: 0.9567
Epoch 20/20
81/81 [==============================] - 55s 683ms/step - loss: 0.0176 - accuracy: 0.9922 - val_loss: 0.2275 - val_accuracy: 0.9603

ผลจากการ Training ทั้งหมด 20 Epochs เห็นได้ว่าในแต่ละ Epoch GPU ใช้เวลาในการประมวลผลที่ราว ๆ 700ms (millisecond) ลองเปลี่ยนมาใช้ TPU โดยที่ Setting เหมือนเดิม เพื่อดูประสิทธิภาพการประมวลผล แต่ก่อนอื่น Coding สำหรับการใช้ TPU มีการเปลี่ยนจากเดิมไปบ้าง ไม่เหมือนกับการทำงานบน CPU และ GPU สามารถใช้งาน Code เดียวกันได้ ดังนั้นต้องมีการปรับแก้ก่อนเริ่ม Training

Code

Code 3 บรรทัดด้านบนเป็นการเปิดการทำงานของ TPU บน Google Colab

การ Training model ด้วย TPU ต้องใช้งานจาก tensorflow.distribute.TPUStrategy() โดยใช้ตัวแปรมารับ 1 ตัว แล้วเริ่มใช้งานตัวแปรนั้น ๆ ด้วย with statement .scope() ซึ่งสามารถออกแบบ Deep learning model ได้อย่างที่เคยทำตามปกติ (ออกแบบเหมือนเดิม)

Epoch 1/20
81/81 [==============================] - 39s 477ms/step - loss: 2.1866 - accuracy: 0.8729 - val_loss: 0.1320 - val_accuracy: 0.9603
Epoch 2/20
81/81 [==============================] - 9s 111ms/step - loss: 0.0849 - accuracy: 0.9736 - val_loss: 0.1075 - val_accuracy: 0.9621
Epoch 3/20
81/81 [==============================] - 9s 111ms/step - loss: 0.0672 - accuracy: 0.9814 - val_loss: 0.0943 - val_accuracy: 0.9657
Epoch 4/20
81/81 [==============================] - 9s 111ms/step - loss: 0.0333 - accuracy: 0.9884 - val_loss: 0.1128 - val_accuracy: 0.9657
Epoch 5/20
81/81 [==============================] - 9s 111ms/step - loss: 0.0267 - accuracy: 0.9907 - val_loss: 0.1152 - val_accuracy: 0.9621
Epoch 6/20
81/81 [==============================] - 9s 112ms/step - loss: 0.5040 - accuracy: 0.9512 - val_loss: 11.1023 - val_accuracy: 0.8610
Epoch 7/20
81/81 [==============================] - 9s 112ms/step - loss: 1.1332 - accuracy: 0.9465 - val_loss: 0.2224 - val_accuracy: 0.9332
Epoch 8/20
81/81 [==============================] - 9s 116ms/step - loss: 0.0685 - accuracy: 0.9775 - val_loss: 0.1721 - val_accuracy: 0.9585
Epoch 9/20
81/81 [==============================] - 9s 112ms/step - loss: 0.0334 - accuracy: 0.9891 - val_loss: 0.2150 - val_accuracy: 0.9603
Epoch 10/20
81/81 [==============================] - 9s 111ms/step - loss: 0.0298 - accuracy: 0.9938 - val_loss: 0.2034 - val_accuracy: 0.9585
Epoch 11/20
81/81 [==============================] - 9s 111ms/step - loss: 0.0174 - accuracy: 0.9984 - val_loss: 0.2282 - val_accuracy: 0.9603
Epoch 12/20
81/81 [==============================] - 9s 111ms/step - loss: 0.0204 - accuracy: 0.9969 - val_loss: 0.2545 - val_accuracy: 0.9585
Epoch 13/20
81/81 [==============================] - 9s 111ms/step - loss: 0.0113 - accuracy: 0.9984 - val_loss: 0.3489 - val_accuracy: 0.9567
Epoch 14/20
81/81 [==============================] - 9s 113ms/step - loss: 0.0212 - accuracy: 0.9961 - val_loss: 0.2691 - val_accuracy: 0.9657
Epoch 15/20
81/81 [==============================] - 9s 112ms/step - loss: 0.0189 - accuracy: 0.9977 - val_loss: 0.2887 - val_accuracy: 0.9585
Epoch 16/20
81/81 [==============================] - 9s 111ms/step - loss: 0.6753 - accuracy: 0.9620 - val_loss: 0.6478 - val_accuracy: 0.9675
Epoch 17/20
81/81 [==============================] - 9s 111ms/step - loss: 0.0912 - accuracy: 0.9806 - val_loss: 0.3511 - val_accuracy: 0.9531
Epoch 18/20
81/81 [==============================] - 9s 112ms/step - loss: 0.0156 - accuracy: 0.9969 - val_loss: 0.3724 - val_accuracy: 0.9585
Epoch 19/20
81/81 [==============================] - 9s 112ms/step - loss: 0.0097 - accuracy: 0.9984 - val_loss: 0.3924 - val_accuracy: 0.9621
Epoch 20/20
81/81 [==============================] - 9s 112ms/step - loss: 0.0210 - accuracy: 0.9946 - val_loss: 0.4639 - val_accuracy: 0.9549

เห็นได้ว่าในแต่ละ Epoch โมเดลถูก Training ด้วยเวลาที่สั้นลง โดยเวลาที่ใช้อยู่ที่ราว ๆ 100 ms ถ้าวัดด้วยความรู้สึกแล้วอาจไม่เห็นความแตกต่างมาก เพราะหน่วยเป็น Millisecond แต่หากคิดเลขเร็ว ๆ จะพบว่าความเร็วในการทำงานต่างกันถึง 6 เท่าตัว ดังนั้นจึงถือว่าเป็นอีกทางเลือกที่ดีสำหรับการทำงานสายนี้