| 36 | |
| 37 | |
| 38 | def train_fn(generator,discriminator,optimizer_G,optimizer_D,adversarial_loss,dataloader,epoch): |
| 39 | loop = tqdm(dataloader,leave=True) |
| 40 | loop.set_description(desc="Training: ") |
| 41 | for step,(imgs,labels) in enumerate(loop): |
| 42 | valid = torch.ones(config.BATCH_SIZE).to(config.DEVICE) |
| 43 | fake = torch.zeros(config.BATCH_SIZE).to(config.DEVICE) |
| 44 | |
| 45 | # print('labels.shape: {}'.format(np.shape(labels))) |
| 46 | # print('labels: {}'.format(labels)) |
| 47 | |
| 48 | real_imgs = imgs.to(config.DEVICE) |
| 49 | real_y = torch.zeros(config.BATCH_SIZE,config.NUM_CLASSES) |
| 50 | #dim = 1,index = config.BATCH_SIZE,src = 1 |
| 51 | #https://mbd.baidu.com/ma/s/HT3QuRvI |
| 52 | real_y = real_y.scatter_(1, labels.view(config.BATCH_SIZE, 1), 1).view(config.BATCH_SIZE,config.NUM_CLASSES,1,1) |
| 53 | real_y = real_y.expand(-1,-1,config.IMG_SIZE,config.IMG_SIZE).to(config.DEVICE) |
| 54 | |
| 55 | noise = torch.randn(size = (config.BATCH_SIZE,config.LATENT_DIM,1,1)).to(config.DEVICE) |
| 56 | gen_labels = (torch.rand(config.BATCH_SIZE,1)*config.NUM_CLASSES).type(torch.LongTensor) |
| 57 | |
| 58 | # print('gen_labels.shape: {}'.format(gen_labels.shape)) => [BATCH_SIZE,1] |
| 59 | # print('gen_labels: {}'.format(gen_labels))[[3.6908],[8.2607],[7.1017],......] |
| 60 | |
| 61 | gen_y = torch.zeros(config.BATCH_SIZE,config.NUM_CLASSES) |
| 62 | # dim = 1,index = config.BATCH_SIZE,src = 1 |
| 63 | gen_y = gen_y.scatter_(1, gen_labels.view(config.BATCH_SIZE, 1), 1).view(config.BATCH_SIZE,config.NUM_CLASSES,1,1) |
| 64 | gen_y = gen_y.to(config.DEVICE) |
| 65 | |
| 66 | #expand :https://blog.csdn.net/weixin_39504171/article/details/106090626/ |
| 67 | gen_y_for_D = gen_y.view(config.BATCH_SIZE, config.NUM_CLASSES, 1, 1).expand(-1, -1, config.IMG_SIZE, config.IMG_SIZE) |
| 68 | gen_y_for_D = gen_y_for_D.to(config.DEVICE) |
| 69 | |
| 70 | #compute the discriminator's loss |
| 71 | optimizer_D.zero_grad() |
| 72 | d_real_loss = adversarial_loss(np.squeeze(discriminator(real_imgs,real_y)),valid) |
| 73 | gen_imgs = generator(noise,gen_y) |
| 74 | d_fake_loss = adversarial_loss(np.squeeze(discriminator(gen_imgs.detach(),gen_y_for_D)),fake) |
| 75 | # d_fake_loss = adversarial_loss(np.squeeze(discriminator(gen_imgs.detach(), real_y)), fake) |
| 76 | d_loss = d_real_loss + d_fake_loss |
| 77 | d_loss.backward() |
| 78 | optimizer_D.step() |
| 79 | |
| 80 | #compute the generator's loss |
| 81 | optimizer_G.zero_grad() |
| 82 | g_loss = adversarial_loss(np.squeeze(discriminator(gen_imgs,gen_y_for_D)),valid) |
| 83 | # g_loss = adversarial_loss(np.squeeze(discriminator(gen_imgs, real_y)), valid) |
| 84 | g_loss.backward() |
| 85 | optimizer_G.step() |
| 86 | |
| 87 | loop.set_postfix( |
| 88 | epoch = epoch, |
| 89 | d_loss = d_loss.item(), |
| 90 | g_loss = g_loss.item() |
| 91 | ) |
| 92 | |
| 93 | |
| 94 | |