()
| 55 | |
| 56 | |
| 57 | def main(): |
| 58 | disc = Discriminator.Discriminator(in_channels=3).to(config.DEVICE) |
| 59 | gen = Generator.Generator(in_channles=3).to(config.DEVICE) |
| 60 | |
| 61 | opt_disc = torch.optim.Adam(disc.parameters(),lr = config.LEARNING_RATE,betas=(0.5,0.999)) |
| 62 | opt_gen = torch.optim.Adam(gen.parameters(),lr = config.LEARNING_RATE,betas=(0.5,0.999)) |
| 63 | |
| 64 | BCE = torch.nn.BCEWithLogitsLoss() |
| 65 | L1_LOSS = torch.nn.L1Loss() |
| 66 | |
| 67 | if config.LOAD_MODEL: |
| 68 | load_checkpoint(config.CHECKPOINT_GEN,gen,opt_gen,config.LEARNING_RATE) |
| 69 | load_checkpoint(config.CHECKPOINT_DISC, gen, opt_disc, config.LEARNING_RATE) |
| 70 | |
| 71 | train_dataset = MapDataset(root_dir=config.TRAIN_DIR) |
| 72 | val_dataset = MapDataset(root_dir=config.VAL_DIR) |
| 73 | |
| 74 | train_loader = DataLoader(dataset=train_dataset,batch_size=config.BATCH_SIZE,shuffle=True) |
| 75 | val_loader = DataLoader(dataset=val_dataset,batch_size=1,shuffle=True) |
| 76 | |
| 77 | # g_scaler = torch.cuda.amp.GradScaler() |
| 78 | # d_scaler = torch.cuda.amp.GradScaler() |
| 79 | |
| 80 | for epoch in range(config.NUM_EPOCHS): |
| 81 | # train_fn(disc,gen,train_loader,opt_disc,opt_gen,L1_LOSS,BCE,g_scaler,d_scaler) |
| 82 | train_fn(disc, gen, train_loader, opt_disc, opt_gen, L1_LOSS, BCE) |
| 83 | |
| 84 | if config.SAVE_MODEL and epoch % 20 == 0: |
| 85 | save_checkpoint(gen,opt_gen,config.CHECKPOINT_GEN) |
| 86 | save_checkpoint(disc,opt_disc,config.CHECKPOINT_DISC) |
| 87 | save_some_examples(gen,val_loader,epoch,folder="images") |
| 88 | |
| 89 | |
| 90 | if __name__ == '__main__': |
no test coverage detected