| 17 | from utils import save_checkpoint,save_some_examples,load_checkpoint |
| 18 | |
| 19 | def train_fn(disc,gen,train_loader,opt_disc,opt_gen,L1_LOSS,BCE): |
| 20 | loop = tqdm(train_loader,leave=True) |
| 21 | for idx,(x,y) in enumerate(loop): |
| 22 | x,y = x.to(config.DEVICE),y.to(config.DEVICE) |
| 23 | |
| 24 | #train dsicriminator |
| 25 | # with torch.cuda.amp.autocast(): x-对应的是卫星拍摄的真实图 y-表示对应卫星拍摄的Google map |
| 26 | y_fake = gen(x) |
| 27 | D_real = disc(x, y) |
| 28 | D_fake = disc(x,y_fake.detach()) |
| 29 | D_real_loss = BCE(D_real,torch.ones_like(D_real)) |
| 30 | D_fake_loss = BCE(D_fake,torch.zeros_like(D_fake)) |
| 31 | D_loss = (D_real_loss + D_fake_loss) / 2 |
| 32 | |
| 33 | opt_disc.zero_grad() |
| 34 | D_loss.backward() |
| 35 | opt_disc.step() |
| 36 | |
| 37 | #train generator |
| 38 | # with torch.cuda.amp.autocast(): |
| 39 | D_fake = disc(x,y_fake) |
| 40 | G_fake_loss = BCE(D_fake,torch.ones_like(D_fake)) |
| 41 | L1 = L1_LOSS(y_fake,y)*config.L1_LAMBDA |
| 42 | G_loss = G_fake_loss + L1 |
| 43 | |
| 44 | opt_gen.zero_grad() |
| 45 | G_loss.backward() |
| 46 | opt_gen.step() |
| 47 | |
| 48 | if idx % 10 == 0: |
| 49 | #设置进度条显示的信息，下面表示在显示过程中同时显示损失值 |
| 50 | loop.set_postfix( |
| 51 | D_real=torch.sigmoid(D_real).mean().item(), |
| 52 | D_fake=torch.sigmoid(D_fake).mean().item(), |
| 53 | ) |
| 54 | |
| 55 | |
| 56 | |