()
| 104 | |
| 105 | |
| 106 | def main_(): |
| 107 | disc_X = Discriminator(in_channels=3).to(config.DEVICE) |
| 108 | disc_Y = Discriminator(in_channels=3).to(config.DEVICE) |
| 109 | gen_F = Generator(img_channels=3,num_residual=9).to(config.DEVICE) |
| 110 | gen_G = Generator(img_channels=3,num_residual=9).to(config.DEVICE) |
| 111 | opt_disc = torch.optim.Adam( |
| 112 | list(disc_X.parameters()) + list(disc_Y.parameters()), |
| 113 | lr = config.LEARNING_RATE, |
| 114 | betas=(0.5,0.999) |
| 115 | ) |
| 116 | opt_gen = torch.optim.Adam( |
| 117 | list(gen_G.parameters()) + list(gen_F.parameters()), |
| 118 | lr = config.LEARNING_RATE, |
| 119 | betas=(0.5,0.999) |
| 120 | ) |
| 121 | L1 = torch.nn.L1Loss() |
| 122 | mse = torch.nn.MSELoss() |
| 123 | |
| 124 | #导入预训练模型 |
| 125 | if config.LOAD_MODEL: |
| 126 | utils.load_checkpoin( |
| 127 | config.CHECKPOINT_GEN_H,gen_G,opt_gen,config.LEARNING_RATE |
| 128 | ) |
| 129 | utils.load_checkpoin( |
| 130 | config.CHECKPOINT_GEN_Z, gen_F, opt_gen, config.LEARNING_RATE |
| 131 | ) |
| 132 | utils.load_checkpoin( |
| 133 | config.CHECKPOINT_CRITICH_H, disc_X, opt_disc, config.LEARNING_RATE |
| 134 | ) |
| 135 | utils.load_checkpoin( |
| 136 | config.CHECKPOINT_CRITICH_Z, disc_Y, opt_disc, config.LEARNING_RATE |
| 137 | ) |
| 138 | |
| 139 | dataset = vanGoPhotoDataset( |
| 140 | root_vango=config.TRAIN_DIR + "/trainA",root_photo=config.TRAIN_DIR + "/trainB", |
| 141 | transform=config.transform |
| 142 | ) |
| 143 | |
| 144 | loader = DataLoader( |
| 145 | dataset=dataset, |
| 146 | batch_size=config.BATCH_SIZE, |
| 147 | shuffle=True, |
| 148 | num_workers=config.NUM_WORKERS, |
| 149 | pin_memory=False |
| 150 | ) |
| 151 | g_scale = 0 |
| 152 | d_scale = 0 |
| 153 | cudaIsAvailable = False |
| 154 | if torch.cuda.is_available(): |
| 155 | g_scale = torch.cuda.amp.GradScaler() |
| 156 | d_scale = torch.cuda.amp.GradScaler() |
| 157 | cudaIsAvailable = True |
| 158 | |
| 159 | for epoch in range(config.NUM_EPOCHS): |
| 160 | train_fn(disc_X,disc_Y,gen_G,gen_F,loader,opt_disc,opt_gen,L1,mse,d_scale,g_scale,epoch,cudaIsAvailable) |
| 161 | if config.SAVE_MODEL: |
| 162 | utils.save_checkpoint(gen_G,opt_gen,filename=config.CHECKPOINT_GEN_H,epochs = epoch) |
| 163 | utils.save_checkpoint(gen_F, opt_gen, filename=config.CHECKPOINT_GEN_Z,epochs = epoch) |
no test coverage detected