import torch import torch.nn.functional as F from models.SwinIR import compute_gradient_penalty from tqdm import tqdm from .utils import get_lr, show_result from .utils_metrics import PSNR, SSIM def fit_one_epoch(writer, G_model_train, D_model_train, G_model, D_model, VGG_feature_model, ResNeSt_model, G_optimizer, D_optimizer, BCEWithLogits_loss, L1_loss, Face_loss, epoch, epoch_size, gen, Epoch, cuda, batch_size, save_interval): G_total_loss = 0 D_total_loss = 0 G_total_PSNR = 0 G_total_SSIM = 0 with tqdm(total=epoch_size,desc=f'Epoch {epoch + 1}/{Epoch}',postfix=dict,mininterval=0.3, ncols=150) as pbar: for iteration, batch in enumerate(gen): if iteration >= epoch_size: break with torch.no_grad(): lr_images, hr_images = batch lr_images, hr_images = torch.from_numpy(lr_images).type(torch.FloatTensor), torch.from_numpy(hr_images).type(torch.FloatTensor) y_real, y_fake = torch.ones(batch_size), torch.zeros(batch_size) if cuda: lr_images, hr_images, y_real, y_fake = lr_images.cuda(), hr_images.cuda(), y_real.cuda(), y_fake.cuda() #-------------------------------------------------# # 训练判别器 #-------------------------------------------------# D_optimizer.zero_grad() D_result_r = D_model_train(hr_images) G_result = G_model_train(lr_images) D_result_f = D_model_train(G_result).squeeze() D_result_rf = D_result_r - D_result_f.mean() D_result_fr = D_result_f - D_result_r.mean() D_train_loss_rf = BCEWithLogits_loss(D_result_rf, y_real) D_train_loss_fr = BCEWithLogits_loss(D_result_fr, y_fake) gradient_penalty = compute_gradient_penalty(D_model_train, hr_images, G_result) D_train_loss = 10 * gradient_penalty + (D_train_loss_rf + D_train_loss_fr) / 2 D_train_loss.backward() D_optimizer.step() #-------------------------------------------------# # 训练生成器 #-------------------------------------------------# G_optimizer.zero_grad() G_result = G_model_train(lr_images) image_loss = L1_loss(G_result, hr_images) D_result_r = D_model_train(hr_images) D_result_f = D_model_train(G_result).squeeze() D_result_rf = D_result_r - D_result_f.mean() D_result_fr = D_result_f - D_result_r.mean() D_train_loss_rf = BCEWithLogits_loss(D_result_rf, y_fake) D_train_loss_fr = BCEWithLogits_loss(D_result_fr, y_real) adversarial_loss = (D_train_loss_rf + D_train_loss_fr) / 2 perception_loss = L1_loss(VGG_feature_model(G_result), VGG_feature_model(hr_images)) # 进行下采样以适配人脸识别网络 G_result_face = F.interpolate(G_result, size=(112, 112), mode='bicubic', align_corners=True) hr_images_face = F.interpolate(hr_images, size=(112, 112), mode='bicubic', align_corners=True) face_loss = torch.mean(1. - Face_loss(ResNeSt_model(G_result_face), ResNeSt_model(hr_images_face))) G_train_loss = 3.0 * image_loss + 1.0 * adversarial_loss + 0.9 * perception_loss + 2.5 * face_loss G_train_loss.backward() G_optimizer.step() G_total_loss += G_train_loss.item() D_total_loss += D_train_loss.item() with torch.no_grad(): G_total_PSNR += PSNR(G_result, hr_images).item() G_total_SSIM += SSIM(G_result, hr_images).item() pbar.set_postfix(**{'G_loss' : G_total_loss / (iteration + 1), 'D_loss' : D_total_loss / (iteration + 1), 'G_PSNR' : G_total_PSNR / (iteration + 1), 'G_SSIM' : G_total_SSIM / (iteration + 1), 'lr' : get_lr(G_optimizer)}) pbar.update(1) if iteration % save_interval == 0: show_result(epoch + 1, G_model_train, lr_images, hr_images) writer.add_scalar('G_loss', G_total_loss / (iteration + 1), epoch + 1) writer.add_scalar('D_loss', D_total_loss / (iteration + 1), epoch + 1) writer.add_scalar('G_PSNR', G_total_PSNR / (iteration + 1), epoch + 1) writer.add_scalar('G_SSIM', G_total_SSIM / (iteration + 1), epoch + 1) writer.add_scalar('lr', get_lr(G_optimizer), epoch + 1) print('Epoch:'+ str(epoch + 1) + '/' + str(Epoch)) print('G Loss: %.4f || D Loss: %.4f ' % (G_total_loss / epoch_size, D_total_loss / epoch_size)) print('Saving state, iter:', str(epoch+1)) # 保存模型权重 torch.save(G_model, 'logs/G_Epoch%d-GLoss%.4f-DLoss%.4f.pth'%((epoch + 1), G_total_loss / epoch_size, D_total_loss / epoch_size)) torch.save(D_model, 'logs/D_Epoch%d-GLoss%.4f-DLoss%.4f.pth'%((epoch + 1), G_total_loss / epoch_size, D_total_loss / epoch_size))