update training documents.

Author: pq-yang

docs/train.md

@@ -7,10 +7,13 @@
 ---
 
 ## Training
+```
+For PyTorch versions >= 1.10, please replace `python -m torch.distributed.launch` in the commands below with `torchrun`.
+```
 
 ### 👾 Stage I - VQGAN
 - Training VQGAN:
-  > python -m torch.distributed.launch --nproc_per_node=8 --master_port=4321 basicsr/train.py -opt options/VQGAN_512_ds32_nearest_stage1.yml --launcher pytorch
+  > python -m torch.distributed.launch --nproc_per_node=gpu_num --master_port=4321 basicsr/train.py -opt options/VQGAN_512_ds32_nearest_stage1.yml --launcher pytorch
 
 - After VQGAN training, you can pre-calculate code sequence for the training dataset to speed up the later training stages:
   > python scripts/generate_latent_gt.py
@@ -19,13 +22,13 @@
 
 ### 🚀 Stage II - CodeFormer (w=0)
 - Training Code Sequence Prediction Module:
-  > python -m torch.distributed.launch --nproc_per_node=8 --master_port=4322 basicsr/train.py -opt options/CodeFormer_stage2.yml --launcher pytorch
+  > python -m torch.distributed.launch --nproc_per_node=gpu_num --master_port=4322 basicsr/train.py -opt options/CodeFormer_stage2.yml --launcher pytorch
 
 - Pre-trained CodeFormer of stage II (`codeformer_stage2.pth`) can be found in the folder of Releases v0.1.0: https://github.com/sczhou/CodeFormer/releases/tag/v0.1.0
 
 ### 🛸 Stage III - CodeFormer (w=1)
 - Training Controllable Module:
-  > python -m torch.distributed.launch --nproc_per_node=8 --master_port=4323 basicsr/train.py -opt options/CodeFormer_stage3.yml --launcher pytorch
+  > python -m torch.distributed.launch --nproc_per_node=gpu_num --master_port=4323 basicsr/train.py -opt options/CodeFormer_stage3.yml --launcher pytorch
 
 - Pre-trained CodeFormer (`codeformer.pth`) can be found in the folder of Releases v0.1.0: https://github.com/sczhou/CodeFormer/releases/tag/v0.1.0
 

docs/train_CN.md

@@ -7,10 +7,13 @@
 ---
 
 ## 训练
+```
+对于PyTorch版本 >= 1.10, 请将下面命令中的`python -m torch.distributed.launch`替换为`torchrun`.
+```
 
 ### 👾 阶段 I - VQGAN
 - 训练VQGAN:
-  > python -m torch.distributed.launch --nproc_per_node=8 --master_port=4321 basicsr/train.py -opt options/VQGAN_512_ds32_nearest_stage1.yml --launcher pytorch
+  > python -m torch.distributed.launch --nproc_per_node=gpu_num --master_port=4321 basicsr/train.py -opt options/VQGAN_512_ds32_nearest_stage1.yml --launcher pytorch
 
 - 训练完VQGAN后，可以通过下面代码预先获得训练数据集的密码本序列，从而加速后面阶段的训练过程:
   > python scripts/generate_latent_gt.py
@@ -19,13 +22,13 @@
 
 ### 🚀 阶段 II - CodeFormer (w=0)
 - 训练密码本训练预测模块:
-  > python -m torch.distributed.launch --nproc_per_node=8 --master_port=4322 basicsr/train.py -opt options/CodeFormer_stage2.yml --launcher pytorch
+  > python -m torch.distributed.launch --nproc_per_node=gpu_num --master_port=4322 basicsr/train.py -opt options/CodeFormer_stage2.yml --launcher pytorch
 
 - 预训练CodeFormer第二阶段模型 (`codeformer_stage2.pth`)可以在Releases v0.1.0文档里下载: https://github.com/sczhou/CodeFormer/releases/tag/v0.1.0
 
 ### 🛸 阶段 III - CodeFormer (w=1)
 - 训练可调模块:
-  > python -m torch.distributed.launch --nproc_per_node=8 --master_port=4323 basicsr/train.py -opt options/CodeFormer_stage3.yml --launcher pytorch
+  > python -m torch.distributed.launch --nproc_per_node=gpu_num --master_port=4323 basicsr/train.py -opt options/CodeFormer_stage3.yml --launcher pytorch
 
 - 预训练CodeFormer模型 (`codeformer.pth`)可以在Releases v0.1.0文档里下载: https://github.com/sczhou/CodeFormer/releases/tag/v0.1.0
 

scripts/download_pretrained_models.py

@@ -6,6 +6,8 @@
 
 
 def download_pretrained_models(method, file_urls):
+    if method == 'CodeFormer_train':
+        method = 'CodeFormer'
     save_path_root = f'./weights/{method}'
     os.makedirs(save_path_root, exist_ok=True)
 
@@ -26,6 +28,12 @@ def download_pretrained_models(method, file_urls):
         'CodeFormer': {
             'codeformer.pth': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth'
         },
+        'CodeFormer_train': {
+            'vqgan_code1024.pth': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/vqgan_code1024.pth',
+            'latent_gt_code1024.pth': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/latent_gt_code1024.pth',
+            'codeformer_stage2.pth': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer_stage2.pth',
+            'codeformer.pth': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth'
+        },
         'facelib': {
             # 'yolov5l-face.pth': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/yolov5l-face.pth',
             'detection_Resnet50_Final.pth': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/detection_Resnet50_Final.pth',