feat: fix save model

Author: Makoq

.DS_Store

@@ -5,7 +5,7 @@
 # 图片尺寸
 image_size = 64
 # 数据load多线程处理数
-workers = 2
+workers = 5
 # 生成器特征数
 ngf = 64
 # 判别器特征数
@@ -16,10 +16,10 @@
 nc=3
 
 # Number of GPUs available. Use 0 for CPU mode.
-ngpu = 0
+ngpu = 1
 
 # Number of training epochs
-num_epochs = 5
+num_epochs = 1
 
 # Learning rate for optimizers
 lr = 0.0002 

apps/.DS_Store

@@ -14,6 +14,8 @@
 from device import device
 from plot import plot_original_dataset,loss_plot,image_show
 from const import ngpu, nz, ngf, nc, ndf, batch_size, workers, image_size
+from generator import Generator 
+from discriminator import Discriminator
 
 if __name__ == '__main__':
     #加载数据
@@ -30,8 +32,9 @@
     #绘制生成图像
     image_show(img_list)
     #保存模型
-    netD.save_model('./models')
-    netG.save_model('./models')
+    Generator.save_model(netG,'./models')
+    Discriminator.save(netD,'./models')
+    
     #保存图片
     os.makedirs('./images', exist_ok=True)
     vutils.save_image(img_list[-1], './images/fake_images.png', normalize=True)

apps/GAN/.DS_Store

@@ -35,6 +35,7 @@ def get_generator(nz, ngf, nc):
     #  to ``mean=0``, ``stdev=0.02``.
     netG.apply(weights_init)
 
+
     return netG