Update README.md, assets/loss.png, train.py files

README.md

@@ -73,17 +73,23 @@ pip install -r requirements.txt
 ```
 ## 数据集
 
-- unconditional models训练集[FFHQ](https://drive.google.com/drive/folders/1u2xu7bSrWxrbUxk-dT-UvEJq8IjdmNTP)，
+- unconditional models训练集[FFHQ](https://drive.google.com/drive/folders/1u2xu7bSrWxrbUxk-dT-UvEJq8IjdmNTP)，SCNet快速下载链接[FFHQ_SCNet](http://113.200.138.88:18080/aidatasets/project-dependency/ffhq_dataset)
 
 - 用于快速复现的小数据集SCNet快速下载链接 [images](http://113.200.138.88:18080/aidatasets/project-dependency/stylegan2_pytorch)
 
-下载后解压到当前目录  
+下载后将所有的压缩包解压到/data/ffhq目录中
 数据目录结构如下：
 ```
-./data/ffhq/images
-    └── 00000.png
-    └── 00001.png
-    └── ...
+./data/ffhq/
+    00000
+        └── 00000.png
+        └── 00001.png
+        └── ...
+    01000
+        └── 01000.png
+        └── 01001.png
+        └── ...   
+    ...
 ```
 
 
@@ -91,7 +97,7 @@ pip install -r requirements.txt
 
 先将数据处理为lmdb格式：  
 ```
-python prepare_data.py --out LMDB ./data/ffhq/images
+python prepare_data.py --out LMDB ./data/ffhq
 ```
 
 ### 单机单卡
@@ -103,6 +109,9 @@ HIP_VISIBLE_DEVICES=2 python train.py LMDB --batch 4
 ```
 HIP_VISIBLE_DEVICES=2,3 python -m torch.distributed.launch --nproc_per_node=2 --master_port=29500 train.py --batch 4 LMDB
 ```
+<div align=center>
+    <img src="./assets/loss.png"/>
+</div>  
 
 ## 推理
 下载权重文件[stylegan2-ffhq-config-f.pt](https://pan.baidu.com/s/1IOD4DjjjrMZBF_TfL7pGJQ?pwd=1234)  

train.py

@@ -337,7 +337,7 @@ if __name__ == "__main__":
     parser.add_argument("path", type=str, help="path to the lmdb dataset")
     parser.add_argument('--arch', type=str, default='stylegan2', help='model architectures (stylegan2 | swagan)')
     parser.add_argument(
-        "--iter", type=int, default=800000, help="total training iterations"
+        "--iter", type=int, default=80, help="total training iterations"
     )
     parser.add_argument(
         "--batch", type=int, default=16, help="batch sizes for each gpus"