import torch import pandas as pd from .OCR_network import * from torch.nn import CTCLoss, MSELoss, L1Loss from torch.nn.utils import clip_grad_norm_ import random import unicodedata import sys import torchvision.models as models from models.transformer import * from .BigGAN_networks import * from params import * from .OCR_network import * from models.blocks import LinearBlock, Conv2dBlock, ResBlocks, ActFirstResBlock from util.util import toggle_grad, loss_hinge_dis, loss_hinge_gen, ortho, default_ortho, toggle_grad, prepare_z_y, \ make_one_hot, to_device, multiple_replace, random_word #from models.inception import InceptionV3, calculate_frechet_distance from data.dataset import TextDataset, TextDatasetval import cv2 import time import matplotlib.pyplot as plt import shutil def get_rgb(x): R = 255 - int(int(x>0.5)*255*(x-0.5)/0.5) G = 0 B = 255 + int(int(x<0.5)*255*(x-0.5)/0.5) return R, G, B def get_page_from_words(word_lists, MAX_IMG_WIDTH = 800): line_all = [] line_t = [] width_t = 0 for i in word_lists: width_t = width_t + i.shape[1] + 16 if width_t>MAX_IMG_WIDTH: line_all.append(np.concatenate(line_t, 1)) line_t = [] width_t = i.shape[1] + 16 line_t.append(i) line_t.append(np.ones((i.shape[0], 16))) if len(line_all) == 0: line_all.append(np.concatenate(line_t, 1)) max_lin_widths = MAX_IMG_WIDTH#max([i.shape[1] for i in line_all]) gap_h = np.ones([16,max_lin_widths]) page_= [] for l in line_all: pad_ = np.ones([l.shape[0],max_lin_widths - l.shape[1]]) page_.append(np.concatenate([l, pad_], 1)) page_.append(gap_h) page = np.concatenate(page_, 0) return page*255 class FCNDecoder(nn.Module): def __init__(self, ups=3, n_res=2, dim=512, out_dim=1, res_norm='adain', activ='relu', pad_type='reflect'): super(FCNDecoder, self).__init__() self.model = [] self.model += [ResBlocks(n_res, dim, res_norm, activ, pad_type=pad_type)] for i in range(ups): self.model += [nn.Upsample(scale_factor=2), Conv2dBlock(dim, dim // 2, 5, 1, 2, norm='in', activation=activ, pad_type=pad_type)] dim //= 2 self.model += [Conv2dBlock(dim, out_dim, 7, 1, 3, norm='none', activation='tanh', pad_type=pad_type)] self.model = nn.Sequential(*self.model) def forward(self, x): y = self.model(x) return y class Generator(nn.Module): def __init__(self): super(Generator, self).__init__() INP_CHANNEL = NUM_EXAMPLES if IS_SEQ: INP_CHANNEL = 1 encoder_layer = TransformerEncoderLayer(TN_HIDDEN_DIM, TN_NHEADS, TN_DIM_FEEDFORWARD, TN_DROPOUT, "relu", True) encoder_norm = nn.LayerNorm(TN_HIDDEN_DIM) if True else None self.encoder = TransformerEncoder(encoder_layer, TN_ENC_LAYERS, encoder_norm) decoder_layer = TransformerDecoderLayer(TN_HIDDEN_DIM, TN_NHEADS, TN_DIM_FEEDFORWARD, TN_DROPOUT, "relu", True) decoder_norm = nn.LayerNorm(TN_HIDDEN_DIM) self.decoder = TransformerDecoder(decoder_layer, TN_DEC_LAYERS, decoder_norm, return_intermediate=True) self.Feat_Encoder = nn.Sequential(*([nn.Conv2d(INP_CHANNEL, 64, kernel_size=7, stride=2, padding=3, bias=False)] +list(models.resnet18(pretrained=True).children())[1:-2])) self.query_embed = nn.Embedding(VOCAB_SIZE, TN_HIDDEN_DIM) self.linear_q = nn.Linear(TN_DIM_FEEDFORWARD, TN_DIM_FEEDFORWARD*8) self.DEC = FCNDecoder(res_norm = 'in') self._muE = nn.Linear(512,512) self._logvarE = nn.Linear(512,512) self._muD = nn.Linear(512,512) self._logvarD = nn.Linear(512,512) self.l1loss = nn.L1Loss() self.noise = torch.distributions.Normal(loc=torch.tensor([0.]), scale=torch.tensor([1.0])) def reparameterize(self, mu, logvar): mu = torch.unbind(mu , 1) logvar = torch.unbind(logvar , 1) outs = [] for m,l in zip(mu, logvar): sigma = torch.exp(l) eps = torch.cuda.FloatTensor(l.size()[0],1).normal_(0,1) eps = eps.expand(sigma.size()) out = m + sigma*eps outs.append(out) return torch.stack(outs, 1) def Eval(self, ST, QRS): batch_size = ST.shape[0] if IS_SEQ: B, N, R, C = ST.shape FEAT_ST = self.Feat_Encoder(ST.view(B*N, 1, R, C)) FEAT_ST = FEAT_ST.view(B, 512, 1, -1) else: FEAT_ST = self.Feat_Encoder(ST) FEAT_ST_ENC = FEAT_ST.flatten(2).permute(2,0,1) memory = self.encoder(FEAT_ST_ENC) if IS_KLD: Ex = memory.permute(1,0,2) memory_mu = self._muE(Ex) memory_logvar = self._logvarE(Ex) memory = self.reparameterize(memory_mu, memory_logvar).permute(1,0,2) OUT_IMGS = [] for i in range(QRS.shape[1]): QR = QRS[:, i, :] if ALL_CHARS: QR_EMB = self.query_embed.weight.repeat(ST.shape[0],1,1).permute(1,0,2) else: QR_EMB = self.query_embed.weight[QR].permute(1,0,2) tgt = torch.zeros_like(QR_EMB) hs = self.decoder(tgt, memory, query_pos=QR_EMB) if IS_KLD: Dx = hs[0].permute(1,0,2) hs_mu = self._muD(Dx) hs_logvar = self._logvarD(Dx) hs = self.reparameterize(hs_mu, hs_logvar).permute(1,0,2).unsqueeze(0) h = hs.transpose(1, 2)[-1]#torch.cat([hs.transpose(1, 2)[-1], QR_EMB.permute(1,0,2)], -1) if ADD_NOISE: h = h + self.noise.sample(h.size()).squeeze(-1).to(DEVICE) h = self.linear_q(h) h = h.contiguous() if ALL_CHARS: h = torch.stack([h[i][QR[i]] for i in range(batch_size)], 0) h = h.view(h.size(0), h.shape[1]*2, 4, -1) h = h.permute(0, 3, 2, 1) h = self.DEC(h) OUT_IMGS.append(h.detach()) return OUT_IMGS def forward(self, ST, QR, QRs = None, mode = 'train'): #Attention Visualization Init enc_attn_weights, dec_attn_weights = [], [] self.hooks = [ self.encoder.layers[-1].self_attn.register_forward_hook( lambda self, input, output: enc_attn_weights.append(output[1]) ), self.decoder.layers[-1].multihead_attn.register_forward_hook( lambda self, input, output: dec_attn_weights.append(output[1]) ), ] #Attention Visualization Init B, N, R, C = ST.shape FEAT_ST = self.Feat_Encoder(ST.view(B*N, 1, R, C)) FEAT_ST = FEAT_ST.view(B, 512, 1, -1) FEAT_ST_ENC = FEAT_ST.flatten(2).permute(2,0,1) memory = self.encoder(FEAT_ST_ENC) QR_EMB = self.query_embed.weight[QR].permute(1,0,2) tgt = torch.zeros_like(QR_EMB) hs = self.decoder(tgt, memory, query_pos=QR_EMB) h = hs.transpose(1, 2)[-1]#torch.cat([hs.transpose(1, 2)[-1], QR_EMB.permute(1,0,2)], -1) if ADD_NOISE: h = h + self.noise.sample(h.size()).squeeze(-1).to(DEVICE) h = self.linear_q(h) h = h.contiguous() h = h.view(h.size(0), h.shape[1]*2, 4, -1) h = h.permute(0, 3, 2, 1) h = self.DEC(h) self.dec_attn_weights = dec_attn_weights[-1].detach() self.enc_attn_weights = enc_attn_weights[-1].detach() for hook in self.hooks: hook.remove() return h class TRGAN(nn.Module): def __init__(self, batch_size=batch_size): super(TRGAN, self).__init__() self.batch_size = batch_size self.epsilon = 1e-7 self.netG = Generator().to(DEVICE) self.netD = nn.DataParallel(Discriminator()).to(DEVICE) self.netW = nn.DataParallel(WDiscriminator()).to(DEVICE) self.netconverter = strLabelConverter(ALPHABET) self.netOCR = CRNN().to(DEVICE) self.OCR_criterion = CTCLoss(zero_infinity=True, reduction='none') #block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[2048] #self.inception = InceptionV3([block_idx]).to(DEVICE) self.optimizer_G = torch.optim.Adam(self.netG.parameters(), lr=G_LR, betas=(0.0, 0.999), weight_decay=0, eps=1e-8) self.optimizer_OCR = torch.optim.Adam(self.netOCR.parameters(), lr=OCR_LR, betas=(0.0, 0.999), weight_decay=0, eps=1e-8) self.optimizer_D = torch.optim.Adam(self.netD.parameters(), lr=D_LR, betas=(0.0, 0.999), weight_decay=0, eps=1e-8) self.optimizer_wl = torch.optim.Adam(self.netW.parameters(), lr=W_LR, betas=(0.0, 0.999), weight_decay=0, eps=1e-8) self.optimizers = [self.optimizer_G, self.optimizer_OCR, self.optimizer_D, self.optimizer_wl] self.optimizer_G.zero_grad() self.optimizer_OCR.zero_grad() self.optimizer_D.zero_grad() self.optimizer_wl.zero_grad() self.loss_G = 0 self.loss_D = 0 self.loss_Dfake = 0 self.loss_Dreal = 0 self.loss_OCR_fake = 0 self.loss_OCR_real = 0 self.loss_w_fake = 0 self.loss_w_real = 0 self.Lcycle1 = 0 self.Lcycle2 = 0 self.lda1 = 0 self.lda2 = 0 self.KLD = 0 with open(ENGLISH_WORDS_PATH, 'rb') as f: self.lex = f.read().splitlines() lex=[] for word in self.lex: try: word=word.decode("utf-8") except: continue if len(word)<20: lex.append(word) self.lex = lex f = open('mytext.txt', 'r') self.text = [j.encode() for j in sum([i.split(' ') for i in f.readlines()], [])]#[:NUM_EXAMPLES] self.eval_text_encode, self.eval_len_text = self.netconverter.encode(self.text) self.eval_text_encode = self.eval_text_encode.to(DEVICE).repeat(self.batch_size, 1, 1) def save_images_for_fid_calculation(self, dataloader, epoch, mode = 'train'): self.real_base = os.path.join('saved_images', EXP_NAME, 'Real') self.fake_base = os.path.join('saved_images', EXP_NAME, 'Fake') if os.path.isdir(self.real_base): shutil.rmtree(self.real_base) if os.path.isdir(self.fake_base): shutil.rmtree(self.fake_base) os.mkdir(self.real_base) os.mkdir(self.fake_base) for step,data in enumerate(dataloader): ST = data['simg'].cuda() self.fakes = self.netG.Eval(ST, self.eval_text_encode) fake_images = torch.cat(self.fakes, 1).detach().cpu().numpy() for i in range(fake_images.shape[0]): for j in range(fake_images.shape[1]): #cv2.imwrite(os.path.join(self.real_base, str(step*batch_size + i)+'_'+str(j)+'.png'), 255*(real_images[i,j])) cv2.imwrite(os.path.join(self.fake_base, str(step*self.batch_size + i)+'_'+str(j)+'.png'), 255*(fake_images[i,j])) if mode == 'train': TextDatasetObj = TextDataset(num_examples = self.eval_text_encode.shape[1]) dataset_real = torch.utils.data.DataLoader( TextDatasetObj, batch_size=self.batch_size, shuffle=True, num_workers=0, pin_memory=True, drop_last=True, collate_fn=TextDatasetObj.collate_fn) elif mode == 'test': TextDatasetObjval = TextDatasetval(num_examples = self.eval_text_encode.shape[1]) dataset_real = torch.utils.data.DataLoader( TextDatasetObjval, batch_size=self.batch_size, shuffle=True, num_workers=0, pin_memory=True, drop_last=True, collate_fn=TextDatasetObjval.collate_fn) for step,data in enumerate(dataset_real): real_images = data['simg'].numpy() for i in range(real_images.shape[0]): for j in range(real_images.shape[1]): cv2.imwrite(os.path.join(self.real_base, str(step*self.batch_size + i)+'_'+str(j)+'.png'), 255*(real_images[i,j])) return self.real_base, self.fake_base def _generate_page(self, ST, SLEN, eval_text_encode = None, eval_len_text = None, no_concat = False): if eval_text_encode == None: eval_text_encode = self.eval_text_encode if eval_len_text == None: eval_len_text = self.eval_len_text self.fakes = self.netG.Eval(ST, eval_text_encode) page1s = [] page2s = [] for batch_idx in range(self.batch_size): word_t = [] word_l = [] gap = np.ones([IMG_HEIGHT,16]) line_wids = [] for idx, fake_ in enumerate(self.fakes): word_t.append((fake_[batch_idx,0,:,:eval_len_text[idx]*resolution].cpu().numpy()+1)/2) word_t.append(gap) if len(word_t) == 16 or idx == len(self.fakes) - 1: line_ = np.concatenate(word_t, -1) word_l.append(line_) line_wids.append(line_.shape[1]) word_t = [] gap_h = np.ones([16,max(line_wids)]) page_= [] for l in word_l: pad_ = np.ones([IMG_HEIGHT,max(line_wids) - l.shape[1]]) page_.append(np.concatenate([l, pad_], 1)) page_.append(gap_h) page1 = np.concatenate(page_, 0) word_t = [] word_l = [] gap = np.ones([IMG_HEIGHT,16]) line_wids = [] sdata_ = [i.unsqueeze(1) for i in torch.unbind(ST, 1)] for idx, st in enumerate((sdata_)): word_t.append((st[batch_idx,0,:,:int(SLEN.cpu().numpy()[batch_idx][idx])].cpu().numpy()+1)/2) word_t.append(gap) if len(word_t) == 16 or idx == len(sdata_) - 1: line_ = np.concatenate(word_t, -1) word_l.append(line_) line_wids.append(line_.shape[1]) word_t = [] gap_h = np.ones([16,max(line_wids)]) page_= [] for l in word_l: pad_ = np.ones([IMG_HEIGHT,max(line_wids) - l.shape[1]]) page_.append(np.concatenate([l, pad_], 1)) page_.append(gap_h) page2 = np.concatenate(page_, 0) merge_w_size = max(page1.shape[0], page2.shape[0]) if page1.shape[0] != merge_w_size: page1 = np.concatenate([page1, np.ones([merge_w_size-page1.shape[0], page1.shape[1]])], 0) if page2.shape[0] != merge_w_size: page2 = np.concatenate([page2, np.ones([merge_w_size-page2.shape[0], page2.shape[1]])], 0) page1s.append(page1) page2s.append(page2) #page = np.concatenate([page2, page1], 1) if no_concat: return page2s, page1s page1s_ = np.concatenate(page1s,0) max_wid = max([i.shape[1] for i in page2s]) padded_page2s = [] for para in page2s: padded_page2s.append(np.concatenate([para, np.ones([ para.shape[0], max_wid-para.shape[1]])], 1)) padded_page2s_ = np.concatenate(padded_page2s,0) return np.concatenate([padded_page2s_, page1s_], 1) def get_current_losses(self): losses = {} losses['G'] = self.loss_G losses['D'] = self.loss_D losses['Dfake'] = self.loss_Dfake losses['Dreal'] = self.loss_Dreal losses['OCR_fake'] = self.loss_OCR_fake losses['OCR_real'] = self.loss_OCR_real losses['w_fake'] = self.loss_w_fake losses['w_real'] = self.loss_w_real losses['cycle1'] = self.Lcycle1 losses['cycle2'] = self.Lcycle2 losses['lda1'] = self.lda1 losses['lda2'] = self.lda2 losses['KLD'] = self.KLD return losses def load_networks(self, epoch): BaseModel.load_networks(self, epoch) if self.opt.single_writer: load_filename = '%s_z.pkl' % (epoch) load_path = os.path.join(self.save_dir, load_filename) self.z = torch.load(load_path) def _set_input(self, input): self.input = input def set_requires_grad(self, nets, requires_grad=False): """Set requies_grad=Fasle for all the networks to avoid unnecessary computations Parameters: nets (network list) -- a list of networks requires_grad (bool) -- whether the networks require gradients or not """ if not isinstance(nets, list): nets = [nets] for net in nets: if net is not None: for param in net.parameters(): param.requires_grad = requires_grad def forward(self): self.real = self.input['img'].to(DEVICE) self.label = self.input['label'] self.sdata = self.input['simg'].to(DEVICE) self.ST_LEN = self.input['swids'] self.text_encode, self.len_text = self.netconverter.encode(self.label) self.one_hot_real = make_one_hot(self.text_encode, self.len_text, VOCAB_SIZE).to(DEVICE).detach() self.text_encode = self.text_encode.to(DEVICE).detach() self.len_text = self.len_text.detach() self.words = [word.encode('utf-8') for word in np.random.choice(self.lex, self.batch_size)] self.text_encode_fake, self.len_text_fake = self.netconverter.encode(self.words) self.text_encode_fake = self.text_encode_fake.to(DEVICE) self.one_hot_fake = make_one_hot(self.text_encode_fake, self.len_text_fake, VOCAB_SIZE).to(DEVICE) self.text_encode_fake_js = [] for _ in range(NUM_WORDS - 1): self.words_j = [word.encode('utf-8') for word in np.random.choice(self.lex, self.batch_size)] self.text_encode_fake_j, self.len_text_fake_j = self.netconverter.encode(self.words_j) self.text_encode_fake_j = self.text_encode_fake_j.to(DEVICE) self.text_encode_fake_js.append(self.text_encode_fake_j) self.fake = self.netG(self.sdata, self.text_encode_fake, self.text_encode_fake_js) def backward_D_OCR(self): pred_real = self.netD(self.real.detach()) pred_fake = self.netD(**{'x': self.fake.detach()}) self.loss_Dreal, self.loss_Dfake = loss_hinge_dis(pred_fake, pred_real, self.len_text_fake.detach(), self.len_text.detach(), True) self.loss_D = self.loss_Dreal + self.loss_Dfake self.pred_real_OCR = self.netOCR(self.real.detach()) preds_size = torch.IntTensor([self.pred_real_OCR.size(0)] * self.batch_size).detach() loss_OCR_real = self.OCR_criterion(self.pred_real_OCR, self.text_encode.detach(), preds_size, self.len_text.detach()) self.loss_OCR_real = torch.mean(loss_OCR_real[~torch.isnan(loss_OCR_real)]) loss_total = self.loss_D + self.loss_OCR_real # backward loss_total.backward() for param in self.netOCR.parameters(): param.grad[param.grad!=param.grad]=0 param.grad[torch.isnan(param.grad)]=0 param.grad[torch.isinf(param.grad)]=0 return loss_total def backward_D_WL(self): # Real pred_real = self.netD(self.real.detach()) pred_fake = self.netD(**{'x': self.fake.detach()}) self.loss_Dreal, self.loss_Dfake = loss_hinge_dis(pred_fake, pred_real, self.len_text_fake.detach(), self.len_text.detach(), True) self.loss_D = self.loss_Dreal + self.loss_Dfake self.loss_w_real = self.netW(self.real.detach(), self.input['wcl'].to(DEVICE)).mean() # total loss loss_total = self.loss_D + self.loss_w_real # backward loss_total.backward() return loss_total def optimize_D_WL(self): self.forward() self.set_requires_grad([self.netD], True) self.set_requires_grad([self.netOCR], False) self.set_requires_grad([self.netW], True) self.optimizer_D.zero_grad() self.optimizer_wl.zero_grad() self.backward_D_WL() def backward_D_OCR_WL(self): # Real if self.real_z_mean is None: pred_real = self.netD(self.real.detach()) else: pred_real = self.netD(**{'x': self.real.detach(), 'z': self.real_z_mean.detach()}) # Fake try: pred_fake = self.netD(**{'x': self.fake.detach(), 'z': self.z.detach()}) except: print('a') # Combined loss self.loss_Dreal, self.loss_Dfake = loss_hinge_dis(pred_fake, pred_real, self.len_text_fake.detach(), self.len_text.detach(), self.opt.mask_loss) self.loss_D = self.loss_Dreal + self.loss_Dfake # OCR loss on real data self.pred_real_OCR = self.netOCR(self.real.detach()) preds_size = torch.IntTensor([self.pred_real_OCR.size(0)] * self.opt.batch_size).detach() loss_OCR_real = self.OCR_criterion(self.pred_real_OCR, self.text_encode.detach(), preds_size, self.len_text.detach()) self.loss_OCR_real = torch.mean(loss_OCR_real[~torch.isnan(loss_OCR_real)]) # total loss self.loss_w_real = self.netW(self.real.detach(), self.wcl) loss_total = self.loss_D + self.loss_OCR_real + self.loss_w_real # backward loss_total.backward() for param in self.netOCR.parameters(): param.grad[param.grad!=param.grad]=0 param.grad[torch.isnan(param.grad)]=0 param.grad[torch.isinf(param.grad)]=0 return loss_total def optimize_D_WL_step(self): self.optimizer_D.step() self.optimizer_wl.step() self.optimizer_D.zero_grad() self.optimizer_wl.zero_grad() def backward_OCR(self): # OCR loss on real data self.pred_real_OCR = self.netOCR(self.real.detach()) preds_size = torch.IntTensor([self.pred_real_OCR.size(0)] * self.opt.batch_size).detach() loss_OCR_real = self.OCR_criterion(self.pred_real_OCR, self.text_encode.detach(), preds_size, self.len_text.detach()) self.loss_OCR_real = torch.mean(loss_OCR_real[~torch.isnan(loss_OCR_real)]) # backward self.loss_OCR_real.backward() for param in self.netOCR.parameters(): param.grad[param.grad!=param.grad]=0 param.grad[torch.isnan(param.grad)]=0 param.grad[torch.isinf(param.grad)]=0 return self.loss_OCR_real def backward_D(self): # Real if self.real_z_mean is None: pred_real = self.netD(self.real.detach()) else: pred_real = self.netD(**{'x': self.real.detach(), 'z': self.real_z_mean.detach()}) pred_fake = self.netD(**{'x': self.fake.detach(), 'z': self.z.detach()}) # Combined loss self.loss_Dreal, self.loss_Dfake = loss_hinge_dis(pred_fake, pred_real, self.len_text_fake.detach(), self.len_text.detach(), self.opt.mask_loss) self.loss_D = self.loss_Dreal + self.loss_Dfake # backward self.loss_D.backward() return self.loss_D def backward_G_only(self): self.gb_alpha = 0.7 #self.Lcycle1 = self.Lcycle1.mean() #self.Lcycle2 = self.Lcycle2.mean() self.loss_G = loss_hinge_gen(self.netD(**{'x': self.fake}), self.len_text_fake.detach(), True).mean() pred_fake_OCR = self.netOCR(self.fake) preds_size = torch.IntTensor([pred_fake_OCR.size(0)] * self.batch_size).detach() loss_OCR_fake = self.OCR_criterion(pred_fake_OCR, self.text_encode_fake.detach(), preds_size, self.len_text_fake.detach()) self.loss_OCR_fake = torch.mean(loss_OCR_fake[~torch.isnan(loss_OCR_fake)]) self.loss_G = self.loss_G + self.Lcycle1 + self.Lcycle2 + self.lda1 + self.lda2 - self.KLD self.loss_T = self.loss_G + self.loss_OCR_fake grad_fake_OCR = torch.autograd.grad(self.loss_OCR_fake, self.fake, retain_graph=True)[0] self.loss_grad_fake_OCR = 10**6*torch.mean(grad_fake_OCR**2) grad_fake_adv = torch.autograd.grad(self.loss_G, self.fake, retain_graph=True)[0] self.loss_grad_fake_adv = 10**6*torch.mean(grad_fake_adv**2) self.loss_T.backward(retain_graph=True) grad_fake_OCR = torch.autograd.grad(self.loss_OCR_fake, self.fake, create_graph=True, retain_graph=True)[0] grad_fake_adv = torch.autograd.grad(self.loss_G, self.fake, create_graph=True, retain_graph=True)[0] a = self.gb_alpha * torch.div(torch.std(grad_fake_adv), self.epsilon+torch.std(grad_fake_OCR)) if a is None: print(self.loss_OCR_fake, self.loss_G, torch.std(grad_fake_adv), torch.std(grad_fake_OCR)) if a>1000 or a<0.0001: print(a) self.loss_OCR_fake = a.detach() * self.loss_OCR_fake self.loss_T = self.loss_G + self.loss_OCR_fake self.loss_T.backward(retain_graph=True) grad_fake_OCR = torch.autograd.grad(self.loss_OCR_fake, self.fake, create_graph=False, retain_graph=True)[0] grad_fake_adv = torch.autograd.grad(self.loss_G, self.fake, create_graph=False, retain_graph=True)[0] self.loss_grad_fake_OCR = 10 ** 6 * torch.mean(grad_fake_OCR ** 2) self.loss_grad_fake_adv = 10 ** 6 * torch.mean(grad_fake_adv ** 2) with torch.no_grad(): self.loss_T.backward() if any(torch.isnan(loss_OCR_fake)) or torch.isnan(self.loss_G): print('loss OCR fake: ', loss_OCR_fake, ' loss_G: ', self.loss_G, ' words: ', self.words) sys.exit() def backward_G_WL(self): self.gb_alpha = 0.7 #self.Lcycle1 = self.Lcycle1.mean() #self.Lcycle2 = self.Lcycle2.mean() self.loss_G = loss_hinge_gen(self.netD(**{'x': self.fake}), self.len_text_fake.detach(), True).mean() self.loss_w_fake = self.netW(self.fake, self.input['wcl'].to(DEVICE)).mean() self.loss_G = self.loss_G + self.Lcycle1 + self.Lcycle2 + self.lda1 + self.lda2 - self.KLD self.loss_T = self.loss_G + self.loss_w_fake self.loss_T.backward(retain_graph=True) grad_fake_WL = torch.autograd.grad(self.loss_w_fake, self.fake, create_graph=True, retain_graph=True)[0] grad_fake_adv = torch.autograd.grad(self.loss_G, self.fake, create_graph=True, retain_graph=True)[0] a = self.gb_alpha * torch.div(torch.std(grad_fake_adv), self.epsilon+torch.std(grad_fake_WL)) if a is None: print(self.loss_w_fake, self.loss_G, torch.std(grad_fake_adv), torch.std(grad_fake_WL)) if a>1000 or a<0.0001: print(a) self.loss_w_fake = a.detach() * self.loss_w_fake self.loss_T = self.loss_G + self.loss_w_fake self.loss_T.backward(retain_graph=True) grad_fake_WL = torch.autograd.grad(self.loss_w_fake, self.fake, create_graph=False, retain_graph=True)[0] grad_fake_adv = torch.autograd.grad(self.loss_G, self.fake, create_graph=False, retain_graph=True)[0] self.loss_grad_fake_WL = 10 ** 6 * torch.mean(grad_fake_WL ** 2) self.loss_grad_fake_adv = 10 ** 6 * torch.mean(grad_fake_adv ** 2) with torch.no_grad(): self.loss_T.backward() def backward_G(self): self.opt.gb_alpha = 0.7 self.loss_G = loss_hinge_gen(self.netD(**{'x': self.fake, 'z': self.z}), self.len_text_fake.detach(), self.opt.mask_loss) # OCR loss on real data pred_fake_OCR = self.netOCR(self.fake) preds_size = torch.IntTensor([pred_fake_OCR.size(0)] * self.opt.batch_size).detach() loss_OCR_fake = self.OCR_criterion(pred_fake_OCR, self.text_encode_fake.detach(), preds_size, self.len_text_fake.detach()) self.loss_OCR_fake = torch.mean(loss_OCR_fake[~torch.isnan(loss_OCR_fake)]) self.loss_w_fake = self.netW(self.fake, self.wcl) #self.loss_OCR_fake = self.loss_OCR_fake + self.loss_w_fake # total loss # l1 = self.params[0]*self.loss_G # l2 = self.params[0]*self.loss_OCR_fake #l3 = self.params[0]*self.loss_w_fake self.loss_G_ = 10*self.loss_G + self.loss_w_fake self.loss_T = self.loss_G_ + self.loss_OCR_fake grad_fake_OCR = torch.autograd.grad(self.loss_OCR_fake, self.fake, retain_graph=True)[0] self.loss_grad_fake_OCR = 10**6*torch.mean(grad_fake_OCR**2) grad_fake_adv = torch.autograd.grad(self.loss_G_, self.fake, retain_graph=True)[0] self.loss_grad_fake_adv = 10**6*torch.mean(grad_fake_adv**2) if not False: self.loss_T.backward(retain_graph=True) grad_fake_OCR = torch.autograd.grad(self.loss_OCR_fake, self.fake, create_graph=True, retain_graph=True)[0] grad_fake_adv = torch.autograd.grad(self.loss_G_, self.fake, create_graph=True, retain_graph=True)[0] #grad_fake_wl = torch.autograd.grad(self.loss_w_fake, self.fake, create_graph=True, retain_graph=True)[0] a = self.opt.gb_alpha * torch.div(torch.std(grad_fake_adv), self.epsilon+torch.std(grad_fake_OCR)) #a0 = self.opt.gb_alpha * torch.div(torch.std(grad_fake_adv), self.epsilon+torch.std(grad_fake_wl)) if a is None: print(self.loss_OCR_fake, self.loss_G_, torch.std(grad_fake_adv), torch.std(grad_fake_OCR)) if a>1000 or a<0.0001: print(a) b = self.opt.gb_alpha * (torch.mean(grad_fake_adv) - torch.div(torch.std(grad_fake_adv), self.epsilon+torch.std(grad_fake_OCR))* torch.mean(grad_fake_OCR)) # self.loss_OCR_fake = a.detach() * self.loss_OCR_fake + b.detach() * torch.sum(self.fake) self.loss_OCR_fake = a.detach() * self.loss_OCR_fake #self.loss_w_fake = a0.detach() * self.loss_w_fake self.loss_T = (1-1*self.opt.onlyOCR)*self.loss_G_ + self.loss_OCR_fake# + self.loss_w_fake self.loss_T.backward(retain_graph=True) grad_fake_OCR = torch.autograd.grad(self.loss_OCR_fake, self.fake, create_graph=False, retain_graph=True)[0] grad_fake_adv = torch.autograd.grad(self.loss_G_, self.fake, create_graph=False, retain_graph=True)[0] self.loss_grad_fake_OCR = 10 ** 6 * torch.mean(grad_fake_OCR ** 2) self.loss_grad_fake_adv = 10 ** 6 * torch.mean(grad_fake_adv ** 2) with torch.no_grad(): self.loss_T.backward() else: self.loss_T.backward() if self.opt.clip_grad > 0: clip_grad_norm_(self.netG.parameters(), self.opt.clip_grad) if any(torch.isnan(loss_OCR_fake)) or torch.isnan(self.loss_G_): print('loss OCR fake: ', loss_OCR_fake, ' loss_G: ', self.loss_G, ' words: ', self.words) sys.exit() def optimize_D_OCR(self): self.forward() self.set_requires_grad([self.netD], True) self.set_requires_grad([self.netOCR], True) self.optimizer_D.zero_grad() #if self.opt.OCR_init in ['glorot', 'xavier', 'ortho', 'N02']: self.optimizer_OCR.zero_grad() self.backward_D_OCR() def optimize_OCR(self): self.forward() self.set_requires_grad([self.netD], False) self.set_requires_grad([self.netOCR], True) if self.opt.OCR_init in ['glorot', 'xavier', 'ortho', 'N02']: self.optimizer_OCR.zero_grad() self.backward_OCR() def optimize_D(self): self.forward() self.set_requires_grad([self.netD], True) self.backward_D() def optimize_D_OCR_step(self): self.optimizer_D.step() self.optimizer_OCR.step() self.optimizer_D.zero_grad() self.optimizer_OCR.zero_grad() def optimize_D_OCR_WL(self): self.forward() self.set_requires_grad([self.netD], True) self.set_requires_grad([self.netOCR], True) self.set_requires_grad([self.netW], True) self.optimizer_D.zero_grad() self.optimizer_wl.zero_grad() if self.opt.OCR_init in ['glorot', 'xavier', 'ortho', 'N02']: self.optimizer_OCR.zero_grad() self.backward_D_OCR_WL() def optimize_D_OCR_WL_step(self): self.optimizer_D.step() if self.opt.OCR_init in ['glorot', 'xavier', 'ortho', 'N02']: self.optimizer_OCR.step() self.optimizer_wl.step() self.optimizer_D.zero_grad() self.optimizer_OCR.zero_grad() self.optimizer_wl.zero_grad() def optimize_D_step(self): self.optimizer_D.step() if any(torch.isnan(self.netD.infer_img.blocks[0][0].conv1.bias)): print('D is nan') sys.exit() self.optimizer_D.zero_grad() def optimize_G(self): self.forward() self.set_requires_grad([self.netD], False) self.set_requires_grad([self.netOCR], False) self.set_requires_grad([self.netW], False) self.backward_G() def optimize_G_WL(self): self.forward() self.set_requires_grad([self.netD], False) self.set_requires_grad([self.netOCR], False) self.set_requires_grad([self.netW], False) self.backward_G_WL() def optimize_G_only(self): self.forward() self.set_requires_grad([self.netD], False) self.set_requires_grad([self.netOCR], False) self.set_requires_grad([self.netW], False) self.backward_G_only() def optimize_G_step(self): self.optimizer_G.step() self.optimizer_G.zero_grad() def optimize_ocr(self): self.set_requires_grad([self.netOCR], True) # OCR loss on real data pred_real_OCR = self.netOCR(self.real) preds_size =torch.IntTensor([pred_real_OCR.size(0)] * self.opt.batch_size).detach() self.loss_OCR_real = self.OCR_criterion(pred_real_OCR, self.text_encode.detach(), preds_size, self.len_text.detach()) self.loss_OCR_real.backward() self.optimizer_OCR.step() def optimize_z(self): self.set_requires_grad([self.z], True) def optimize_parameters(self): self.forward() self.set_requires_grad([self.netD], False) self.optimizer_G.zero_grad() self.backward_G() self.optimizer_G.step() self.set_requires_grad([self.netD], True) self.optimizer_D.zero_grad() self.backward_D() self.optimizer_D.step() def test(self): self.visual_names = ['fake'] self.netG.eval() with torch.no_grad(): self.forward() def train_GD(self): self.netG.train() self.netD.train() self.optimizer_G.zero_grad() self.optimizer_D.zero_grad() # How many chunks to split x and y into? x = torch.split(self.real, self.opt.batch_size) y = torch.split(self.label, self.opt.batch_size) counter = 0 # Optionally toggle D and G's "require_grad" if self.opt.toggle_grads: toggle_grad(self.netD, True) toggle_grad(self.netG, False) for step_index in range(self.opt.num_critic_train): self.optimizer_D.zero_grad() with torch.set_grad_enabled(False): self.forward() D_input = torch.cat([self.fake, x[counter]], 0) if x is not None else self.fake D_class = torch.cat([self.label_fake, y[counter]], 0) if y[counter] is not None else y[counter] # Get Discriminator output D_out = self.netD(D_input, D_class) if x is not None: pred_fake, pred_real = torch.split(D_out, [self.fake.shape[0], x[counter].shape[0]]) # D_fake, D_real else: pred_fake = D_out # Combined loss self.loss_Dreal, self.loss_Dfake = loss_hinge_dis(pred_fake, pred_real, self.len_text_fake.detach(), self.len_text.detach(), self.opt.mask_loss) self.loss_D = self.loss_Dreal + self.loss_Dfake self.loss_D.backward() counter += 1 self.optimizer_D.step() # Optionally toggle D and G's "require_grad" if self.opt.toggle_grads: toggle_grad(self.netD, False) toggle_grad(self.netG, True) # Zero G's gradients by default before training G, for safety self.optimizer_G.zero_grad() self.forward() self.loss_G = loss_hinge_gen(self.netD(self.fake, self.label_fake), self.len_text_fake.detach(), self.opt.mask_loss) self.loss_G.backward() self.optimizer_G.step() def save_networks(self, epoch, save_dir): """Save all the networks to the disk. Parameters: epoch (int) -- current epoch; used in the file name '%s_net_%s.pth' % (epoch, name) """ for name in self.model_names: if isinstance(name, str): save_filename = '%s_net_%s.pth' % (epoch, name) save_path = os.path.join(save_dir, save_filename) net = getattr(self, 'net' + name) if len(self.gpu_ids) > 0 and torch.cuda.is_available(): # torch.save(net.module.cpu().state_dict(), save_path) if len(self.gpu_ids) > 1: torch.save(net.module.cpu().state_dict(), save_path) else: torch.save(net.cpu().state_dict(), save_path) net.cuda(self.gpu_ids[0]) else: torch.save(net.cpu().state_dict(), save_path)