BigGAN PyTorch

Biggan-Pytorch

การใช้งาน Pytorch Biggan อย่างเป็นทางการของผู้เขียนอย่างเป็นทางการ

repo นี้มีรหัสสำหรับการฝึกอบรม GPU 4-8 ของ Biggans จากการฝึกอบรม Gan ขนาดใหญ่สำหรับการสังเคราะห์ภาพธรรมชาติที่มีความเที่ยงตรงสูงโดย Andrew Brock, Jeff Donahue และ Karen Simonyan

รหัสนี้คือโดย Andy Brock และ Alex Andonian

วิธีใช้รหัสนี้

คุณจะต้อง:

• Pytorch, เวอร์ชัน 1.0.1
• tqdm, numpy, scipy และ h5py
• ชุดฝึกอบรม Imagenet

ก่อนอื่นคุณอาจเลือกเตรียมชุดข้อมูลเป้าหมาย HDF5 ที่ผ่านการประมวลผลไว้ล่วงหน้าสำหรับ I/O ที่เร็วขึ้น การติดตามสิ่งนี้ (หรือไม่) คุณจะต้องมีช่วงเวลาเริ่มต้นที่จำเป็นในการคำนวณ FID สิ่งเหล่านี้สามารถทำได้โดยการปรับเปลี่ยนและทำงาน

sh scripts/utils/prepare_data.sh

ซึ่งโดยค่าเริ่มต้นจะถือว่าชุดการฝึกอบรม Imagenet ของคุณจะถูกดาวน์โหลดลงใน data โฟลเดอร์รูทในไดเรกทอรีนี้และจะเตรียม HDF5 ที่แคชไว้ที่ความละเอียด 128x128 พิกเซล

ในโฟลเดอร์สคริปต์มีสคริปต์ Bash หลายตัวซึ่งจะฝึกอบรม Biggans ที่มีขนาดแบทช์ที่แตกต่างกัน รหัสนี้จะถือว่าคุณไม่สามารถเข้าถึงฝัก TPU เต็มรูปแบบและการปลอมแปลงแบทช์ขนาดใหญ่โดยใช้การสะสมการไล่ระดับสี โดยค่าเริ่มต้นสคริปต์ launch_BigGAN_bs256x8.sh ฝึกอบรมโมเดล Biggan ขนาดเต็มขนาดเต็มขนาด 256 และ 8 การสะสมการไล่ระดับสีสำหรับขนาดแบทช์ทั้งหมด 2048 ใน 8xv100 พร้อมการฝึกอบรมเต็มความแม่นยำ (ไม่มีแกนเทนเซอร์)

ก่อนอื่นคุณจะต้องหาขนาดแบทช์สูงสุดการตั้งค่าของคุณสามารถรองรับได้ แบบจำลองที่ผ่านการฝึกอบรมมาก่อนที่จะได้รับการฝึกฝนใน 8xv100 (16GB VRAM แต่ละอัน) ซึ่งสามารถรองรับได้มากกว่า BS256 ที่ใช้โดยค่าเริ่มต้นเล็กน้อย เมื่อคุณพิจารณาสิ่งนี้แล้วคุณควรแก้ไขสคริปต์เพื่อให้ขนาดแบทช์คูณจำนวนการสะสมการไล่ระดับสีเท่ากับขนาดแบทช์ทั้งหมดที่คุณต้องการ (ค่าเริ่มต้น Biggan ถึง 2048)

โปรดทราบด้วยว่าสคริปต์นี้ใช้ --load_in_mem arg ซึ่งโหลดไฟล์ทั้งหมด (~ 64GB) i128.hdf5 ทั้งหมดลงใน RAM เพื่อโหลดข้อมูลที่เร็วขึ้น หากคุณมี RAM ไม่เพียงพอที่จะสนับสนุนสิ่งนี้ (อาจเป็น 96GB+) ให้ลบอาร์กิวเมนต์นี้ออก

ตัวชี้วัดและการสุ่มตัวอย่าง

ในระหว่างการฝึกอบรมสคริปต์นี้จะส่งออกบันทึกด้วยตัวชี้วัดการฝึกอบรมและตัวชี้วัดการทดสอบจะช่วยประหยัดสำเนาหลายชุด (2 ครั้งล่าสุดและ 5 คะแนนสูงสุด) ของพารามิเตอร์น้ำหนัก/เครื่องมือเพิ่มประสิทธิภาพและจะผลิตตัวอย่างและการแก้ไขทุกครั้งที่ประหยัดน้ำหนัก โฟลเดอร์บันทึกมีสคริปต์เพื่อประมวลผลบันทึกเหล่านี้และพล็อตผลลัพธ์โดยใช้ MATLAB (ขออภัยไม่ขอโทษ)

หลังจากการฝึกอบรมเราสามารถใช้ sample.py เพื่อสร้างตัวอย่างเพิ่มเติมและการแก้ไขทดสอบด้วยค่าการตัดทอนที่แตกต่างกันขนาดแบทช์จำนวนการสะสมสถิติยืน ฯลฯ ดูตัวอย่างสคริปต์ sample_BigGAN_bs256x8.sh สำหรับตัวอย่าง

โดยค่าเริ่มต้นทุกอย่างจะถูกบันทึกไว้ในน้ำหนัก/ตัวอย่าง/บันทึก/โฟลเดอร์ข้อมูลซึ่งสันนิษฐานว่าอยู่ในโฟลเดอร์เดียวกับ repo นี้ คุณสามารถชี้สิ่งเหล่านี้ทั้งหมดไปยังโฟลเดอร์พื้นฐานที่แตกต่างกันโดยใช้อาร์กิวเมนต์ --base_root หรือเลือกตำแหน่งเฉพาะสำหรับแต่ละรายการด้วยอาร์กิวเมนต์ที่เกี่ยวข้อง (เช่น --logs_root )

เรารวมสคริปต์เพื่อเรียกใช้ Biggan-Deep แต่เรายังไม่ได้ฝึกฝนแบบจำลองที่ใช้พวกเขาอย่างเต็มที่ดังนั้นให้พิจารณาพวกเขาที่ยังไม่ทดลอง นอกจากนี้เรายังรวมสคริปต์เพื่อเรียกใช้โมเดลบน CIFAR และเพื่อเรียกใช้ Sa-Gan (กับ EMA) และ SN-GAN บน Imagenet รหัส SA-GAN จะถือว่าคุณมี 4xtitanx (หรือเทียบเท่าในแง่ของ GPU RAM) และจะทำงานด้วยขนาดแบทช์ 128 และ 2 การสะสมการไล่ระดับสี

หมายเหตุสำคัญเกี่ยวกับตัวชี้วัดการลงทะเบียนเรียน

repo นี้ใช้เครือข่าย Pytorch In Inception ที่สร้างขึ้นเพื่อคำนวณคือและ FID คะแนนเหล่านี้แตกต่างจากคะแนนที่คุณจะได้รับการใช้รหัสการลงทะเบียน TF อย่างเป็นทางการและมีไว้เพื่อการตรวจสอบเท่านั้น! เรียกใช้ sample.py บนโมเดลของคุณด้วยอาร์กิวเมนต์ --sample_npz จากนั้นเรียกใช้ inception_tf13 เพื่อคำนวณ tensorflow จริงคือ โปรดทราบว่าคุณจะต้องมี TensorFlow 1.3 หรือติดตั้งก่อนหน้านี้เนื่องจาก TF1.4+ Breaks ต้นฉบับคือรหัส

นางแบบที่ได้รับการฝึกฝน

เรารวมจุดตรวจสอบแบบจำลองสองแบบ (พร้อม G, D, สำเนา EMA ของ G, Optimizers และ State DICT):

• จุดตรวจสอบหลักสำหรับ Biggan ที่ได้รับการฝึกฝนเกี่ยวกับ Imagenet ที่ 128x128 โดยใช้การสะสม BS256 และ 8 การไล่ระดับ
• จุดตรวจสอบก่อนหน้าของรุ่นแรก (100k g iters) ที่ประสิทธิภาพสูง แต่ก่อนที่จะล่มสลายซึ่งอาจง่ายกว่าที่จะปรับแต่ง: ลิงก์

แบบจำลองที่ได้รับการฝึกฝนสำหรับ Places-365 เร็ว ๆ นี้

repo นี้ยังมีสคริปต์สำหรับการพอร์ตน้ำหนักตัวกำเนิด TFHUB Biggan ดั้งเดิมไปยัง Pytorch ดูสคริปต์ในโฟลเดอร์ TFHUB สำหรับรายละเอียดเพิ่มเติม

ปรับแต่งการใช้ชุดข้อมูลของคุณเองหรือทำฟังก์ชั่นการฝึกอบรมใหม่

หากคุณต้องการกลับมาฝึกซ้อมที่ถูกขัดจังหวะหรือปรับแต่งโมเดลที่ผ่านการฝึกอบรมมาก่อนให้เรียกใช้สคริปต์การเปิดตัวเดียวกัน แต่ด้วยการเพิ่มอาร์กิวเมนต์ --resume ความสัมพันธ์ ชื่อการทดลองจะถูกสร้างขึ้นโดยอัตโนมัติจากการกำหนดค่า แต่สามารถแทนที่ได้โดยใช้ --experiment_name arg (ตัวอย่างเช่นหากคุณต้องการปรับแต่งโมเดลโดยใช้การตั้งค่าเครื่องมือเพิ่มประสิทธิภาพที่แก้ไขแล้ว)

ในการเตรียมชุดข้อมูลของคุณเองคุณจะต้องเพิ่มลงใน Datasets.py และแก้ไขคำสั่งความสะดวกสบายใน utils.py (dset_dict, imsize_dict, root_dict, nclass_dict, classes_per_sheet_dict) เพื่อให้มีเมตาดาต้าที่เหมาะสมสำหรับข้อมูลของคุณ ทำซ้ำกระบวนการใน PREPAL_DATA.SH (เป็นทางเลือกผลิตสำเนาที่ประมวลผลล่วงหน้า HDF5 และคำนวณช่วงเวลาเริ่มต้นสำหรับ FID)

โดยค่าเริ่มต้นสคริปต์การฝึกอบรมจะบันทึกจุดตรวจที่ดีที่สุด 5 อันดับแรกซึ่งวัดจากคะแนนเริ่มต้น สำหรับชุดข้อมูลอื่นนอกเหนือจาก Imagenet คะแนนการลงทะเบียนอาจเป็นตัวชี้วัดคุณภาพที่แย่มากดังนั้นคุณอาจต้องการใช้ --which_best FID แทน

หากต้องการใช้ฟังก์ชั่นการฝึกอบรมของคุณเอง (เช่นฝึก Bigvae): แก้ไข train_fns.gan_training_function หรือเพิ่มรถไฟขบวนใหม่ FN และเพิ่มหลังจาก if config['which_train_fn'] == 'GAN': line in train.py

สิ่งที่เรียบร้อย

• เรารวมบันทึกการฝึกอบรมเต็มรูปแบบและตัวชี้วัดที่นี่เพื่อการอ้างอิง ฉันพบว่าหนึ่งในสิ่งที่ยากที่สุดเกี่ยวกับการปรับปรุงกระดาษอีกครั้งสามารถตรวจสอบได้ว่าบันทึกการจัดเรียงในช่วงต้นของการฝึกอบรมโดยเฉพาะอย่างยิ่งหากการฝึกอบรมใช้เวลาหลายสัปดาห์ หวังว่าสิ่งเหล่านี้จะเป็นประโยชน์สำหรับการทำงานในอนาคต
• เรารวมถึงการคำนวณ FID แบบเร่งความเร็ว-เวอร์ชัน Scipy ดั้งเดิมอาจต้องใช้เวลามากกว่า 10 นาทีในการคำนวณเมทริกซ์ SQRT รุ่นนี้ใช้เวอร์ชัน Pytorch เร่งความเร็วเพื่อคำนวณภายในไม่กี่วินาที
• เรารวมถึงการใช้งานการบริโภค Ortho Reg ที่มีความจำต่ำ
• โดยค่าเริ่มต้นเราจะคำนวณค่าเอกพจน์ด้านบนเท่านั้น (บรรทัดฐานสเปกตรัม) แต่รหัสนี้รองรับการคำนวณ SVS เพิ่มเติมผ่านอาร์กิวเมนต์ --num_G_SVs

ความแตกต่างที่สำคัญระหว่างรหัสนี้และ biggan ดั้งเดิม

• เราใช้การตั้งค่าเครื่องมือเพิ่มประสิทธิภาพจาก SA-GAN (G_LR = 1E-4, D_LR = 4E-4, NUM_D_STEPS = 1 ซึ่งตรงข้ามกับ G_LR ของ Biggan = 5E-5, D_LR = 2E-4, NUM_D_STEPS = 2) ในขณะที่นักแสดงน้อยลงเล็กน้อยนี่เป็นมุมแรกที่เราตัดเพื่อลดเวลาการฝึกซ้อมลง
• โดยค่าเริ่มต้นเราไม่ได้ใช้ cross-replica batchnorm (aka synced batchnorm) ตัวแปรทั้งสองที่เราลอง (แบบกำหนดเองไร้เดียงสาและอันที่รวมอยู่ใน repo นี้) มีการไล่ระดับสีที่แตกต่างกันเล็กน้อย (แม้ว่าจะผ่านไปข้างหน้าเหมือนกัน) จาก batchnorm ในตัวซึ่งดูเหมือนจะเพียงพอที่จะฝึกฝน
• การสะสมการไล่ระดับสีหมายความว่าเราอัปเดตการประมาณการ SV และสถิติ BN บ่อยกว่า 8 เท่า ซึ่งหมายความว่าสถิติ BN นั้นใกล้เคียงกับสถิติที่ยืนอยู่มากและการประมาณค่าเอกพจน์มักจะแม่นยำยิ่งขึ้น ด้วยเหตุนี้เราจึงวัดตัวชี้วัดโดยค่าเริ่มต้นด้วย g ในโหมดทดสอบ (โดยใช้ batchnorm ที่ใช้การประมาณค่าสถิติแทนการคำนวณสถิติการยืนเช่นเดียวกับในกระดาษ) เรายังคงสนับสนุนสถิติยืน (ดูตัวอย่างสคริปต์ shen) สิ่งนี้อาจส่งผลให้เกิดการไล่ระดับสีจากการสะสมก่อนหน้านี้ที่เก่าแก่ แต่ในทางปฏิบัติสิ่งนี้ดูเหมือนจะไม่เป็นปัญหา
• แบบจำลองที่มีอยู่ในปัจจุบันไม่ได้รับการฝึกฝนด้วยการทำให้เป็นมาตรฐานแบบมุมฉาก การฝึกอบรมที่ไม่มี Ortho Reg ดูเหมือนจะเพิ่มความน่าจะเป็นที่โมเดลจะไม่ตอบสนองต่อการตัดทอน แต่ดูเหมือนว่ารุ่นนี้จะได้รับตั๋วที่ชนะ ไม่ว่าเราจะให้การใช้งานที่ดีที่สุด (การใช้หน่วยความจำที่รวดเร็วและน้อยที่สุด) การใช้งาน Ortho Reg ซึ่งคำนวณ Ortho Reg โดยตรง การไล่ระดับสี

หมายเหตุเกี่ยวกับการออกแบบ repo นี้

รหัสนี้ได้รับการออกแบบจากพื้นดินเพื่อใช้เป็นฐานที่ขยายได้และแฮ็กได้สำหรับรหัสการวิจัยเพิ่มเติม เราได้คิดมากเป็นจำนวนมากในการทำให้แน่ใจว่านามธรรมนั้นมีความหนา ที่เหมาะสม สำหรับการวิจัย-ไม่หนาพอที่จะยอมรับไม่ได้ แต่ไม่ผอมเท่าที่ไร้ประโยชน์ แนวคิดหลักคือถ้าคุณต้องการทดสอบกับการตั้งค่า SOTA และทำการปรับเปลี่ยน (ลองใช้ฟังก์ชั่นการสูญเสียใหม่สถาปัตยกรรมบล็อกการใส่ใจในตัวเอง ฯลฯ ) คุณควรทำอย่างง่ายดายเพียงแค่วางรหัสของคุณในหนึ่งหรือสองสถานที่โดยไม่ต้องกังวลเกี่ยวกับส่วนที่เหลือของ codebase สิ่งต่าง ๆ เช่นการใช้ตัวเอง which_conv และ functools.partial ในคำจำกัดความโมเดล biggan.py ถูกนำมารวมกันในใจเช่นเดียวกับการออกแบบของการสืบทอดระดับสเปกตรัมระดับ

ด้วยที่กล่าวว่านี่เป็นรหัสฐานที่ค่อนข้างใหญ่สำหรับโครงการเดียว ในขณะที่เราพยายามอย่างถี่ถ้วนกับความคิดเห็นหากมีบางสิ่งที่คุณคิดว่าอาจชัดเจนขึ้นเขียนดีขึ้นหรือ refactored ดีกว่าโปรดอย่าลังเลที่จะยกปัญหาหรือคำขอดึง

คำขอคุณสมบัติ

ต้องการทำงานหรือปรับปรุงรหัสนี้? มีสองสิ่งที่ repo นี้จะได้รับประโยชน์จาก แต่ยังไม่ได้ผล

• batchnorm ที่ซิงโครไนซ์ (aka cross-replica batchnorm) เราลองใช้สองตัวแปรของสิ่งนี้ แต่ด้วยเหตุผลบางอย่างที่ไม่ทราบสาเหตุการฝึกฝนทำให้พิการในแต่ละครั้ง เรายังไม่ได้ลองใช้ Apex Syncbn เนื่องจากเซิร์ฟเวอร์ของโรงเรียนของฉันอยู่ในไดรเวอร์ Nvidia โบราณที่ไม่สนับสนุน-APEX อาจเป็นจุดเริ่มต้นที่ดี
• การฝึกอบรมที่แม่นยำแบบผสมและใช้ประโยชน์จากแกนเทนเซอร์ repo นี้รวมถึงการใช้งานอดัมที่มีความแม่นยำผสมไร้เดียงสาซึ่งทำงานในช่วงต้นของการฝึกอบรม แต่นำไปสู่การล่มสลายในช่วงต้นและไม่ได้ทำอะไรเพื่อเปิดใช้งานแกนเทนเซอร์ (มันลดการใช้หน่วยความจำ) ดังกล่าวข้างต้นการรวม APEX เข้ากับรหัสนี้และใช้เทคนิคการฝึกอบรมที่มีความแม่นยำแบบผสมเพื่อใช้ประโยชน์จากแกนเทนเซอร์และลดการใช้หน่วยความจำอาจทำให้เกิดความเร็วได้อย่างมาก

หมายเหตุอื่น ๆ

ดูไดเรกทอรีนี้สำหรับฉลาก Imagenet

หากคุณใช้รหัสนี้โปรดอ้างอิง

 @inproceedings{
brock2018large,
title={Large Scale {GAN} Training for High Fidelity Natural Image Synthesis},
author={Andrew Brock and Jeff Donahue and Karen Simonyan},
booktitle={International Conference on Learning Representations},
year={2019},
url={https://openreview.net/forum?id=B1xsqj09Fm},
}

กิตติกรรมประกาศ

ขอบคุณ Google สำหรับการบริจาคเครดิตคลาวด์ที่ใจกว้าง

Syncbn โดย Jiayuan Mao และ Tete Xiao

แถบความคืบหน้ามาจาก Jan Schlüter

Test Metrics Logger จาก VoxNet

การใช้ Pytorch ของ COV จาก Modar M. Alfadly

Pytorch Fast Matrix SQRT สำหรับ FID จาก Tsung-yu Lin และ Subhransu Maji

รหัสคะแนนการลงทะเบียนเรียน Tensorflow จาก Openai ที่ได้รับการปรับปรุง-GAN