def evalRoboustness(net, batch_generator):
defense_accs = AvgMeter()
epsilons = [4, 8, 12, 16, 20, 24]
nb_iters = [40, 80, 120]
Attacks = []
for e in epsilons:
e = e / 255.0
for nb in nb_iters:
Attacks.append(IPGD(e, e//2, nb))
net.eval()
pbar = tqdm(batch_generator)
for data, label in pbar:
data = data.cuda()
label = label.cuda()
choices = np.random.randint(low = 0, high = 17, size = 4)
for c in choices:
defense_accs.update(Attacks[c].get_batch_accuracy(net, data, label))
pbar.set_description('Evulating Roboustness')
return defense_accs.mean
def test(TaskLoader, M, SVGD, DEVICE, num_of_step=3, step_size=1e-3):
raw_M = M
LogP = AvgMeter()
pbar = tqdm(range(100))
for i in pbar:
task = next(TaskLoader)
X, Y, Xtest, Ytest, std = task
X = X.to(DEVICE)
Y = Y.to(DEVICE)
Xtest = Xtest.to(DEVICE)
Ytest = Ytest.to(DEVICE)
std = std.to(DEVICE) * 100 # * 100 to stablize
SVGD.NablaLogP.update(X, Y, std)
SVGD.InitMomentumUpdaters()
#Mt = SVGD.step(M, retain_graph=False, step_size=step_size)
for i in range(num_of_step):
M = SVGD.step(M, retain_graph=False, step_size=step_size)
SVGD.NablaLogP.update(Xtest, Ytest, std)
SVGD.InitMomentumUpdaters()
with torch.no_grad():
logp = 0
for paramsvec in M:
logp = logp + SVGD.NablaLogP(True, paramsvec, ret_grad=False)
logp = logp / len(M)
LogP.update(logp.item())
pbar.set_description("Running Validation")
pbar.set_postfix({'Logp_test': LogP.mean})
M = raw_M
return LogP.mean
def train_epoch(net, loader, optimizer, cost):
# we transfer the mode of network to train
net.train()
batch_loss = AvgMeter()
for batch_idx, (data, label) in enumerate(loader):
data = Variable(
data.cuda()
) # A Variable wraps a Tensor. It supports nearly all the API’s defined by a Tensor.
label = Variable(label.cuda())
output = net(data) # Give the data to the network
loss = cost(output, label)
# evaluate the cost function
output = output.squeeze().data.cpu().numpy()
label = label.squeeze().cpu().numpy()
dice = dice_coeff(output, label)
optimizer.zero_grad(
) # we need to set the gradients to zero before starting to do backpropragation because PyTorch accumulates the gradients on subsequent backward passes
loss.backward()
optimizer.step()
batch_loss.update(loss.item())
if batch_idx % 10 == 0:
print("Train Batch {} || Loss: {:.4f} | Training Dice: {:.4f}".
format(str(batch_idx).zfill(4), batch_loss.val, dice))
return batch_loss.avg
def train_epoch(net, loader, optimizer, cost):
net.train()
labels, predicts = [], []
batch_acc = []
loss_meter = AvgMeter()
t_1 = time.time()
for batch_idx, (feature, label) in enumerate(loader):
feature = feature.float().cuda() #GPU
label = label.float().cuda() #GPU
# feature = feature.float() #CPU
# label = label.float() #CPU
optimizer.zero_grad()
predict = net(feature)
loss = cost(predict, label.long())
loss.backward()
optimizer.step()
loss_meter.update(loss.item())
predict = predict.data.cpu().numpy()
label = label.data.cpu().numpy()
predicts.extend(list(np.argmax(predict, axis=1)))
labels.extend(list(label))
acc = accuracy_score(labels, predicts)
batch_acc.append(acc)
if batch_idx % 10 == 0:
info = [batch_idx, loss_meter.val, acc, time.time() - t_1]
t_1 = time.time()
print("Batch: {} Loss: {:.4f} Batch_acc: {:0.4f} Time:{:0.2f}\n".
format(*info),
end="")
t_acc = np.array(batch_acc)
return loss_meter.avg, acc
def eval(model, loader):
tqdm_object = tqdm(loader, total=len(loader))
loss_meter = AvgMeter()
all_preds = []
for batch in tqdm_object:
batch = {k: v.to(config.DEVICE) for k, v in batch.items()}
preds, loss = model(images=batch['images'], targets=batch['targets'])
all_preds.append(preds)
loss_meter.update(loss.item(), count=batch['images'].size(0))
tqdm_object.set_postfix(loss=loss_meter.avg)
return all_preds, loss_meter
def salience_val_one_epoch(net_f,
net,
optimizer,
batch_generator,
criterion,
SalenceGenerator,
AttackMethod,
clock,
attack_freq=5,
use_adv=True,
DEVICE=torch.device('cuda:0')):
'''
Using net_f to generate salience maps which is used to train a new clsiffier
:param net_f: feature network
:param net:
:param optimizer:
:param batch_generator:
:param criterion:
:param SalenceGenerator:
:param clock:
:return:
'''
clean_acc = AvgMeter()
adv_acc = AvgMeter()
net_f.eval()
net.eval()
#clock.tock()
pbar = tqdm(batch_generator)
start_time = time.time()
for (data, label) in pbar:
#clock.tick()
data = data.to(DEVICE)
label = label.to(DEVICE)
if clock.minibatch % (attack_freq + 1) == 1 and use_adv:
net.eval()
adv_inp = AttackMethod.attack(net_f, data, label)
salience_data = SalenceGenerator(net_f, adv_inp, label)
pred = net(salience_data)
loss = criterion(pred, label)
acc = torch_accuracy(pred, label, (1, ))
adv_acc.update(acc[0].item())
else:
salience_data = SalenceGenerator(net_f, data, label)
pred = net(salience_data)
loss = criterion(pred, label)
acc = torch_accuracy(pred, label, (1, ))
clean_acc.update(acc[0].item())
pbar.set_description("Validation Epoch: {}".format(clock.epoch))
pbar.set_postfix({
'clean_acc': clean_acc.mean,
'adv_acc': adv_acc.mean
})
return {'clean_acc': clean_acc.mean, 'adv_acc': adv_acc.mean}
def train(self, dataset):
for epoch in range(1, config.epochs + 1):
log.info("Epoch : %d", epoch)
pi_losses = AvgMeter()
v_losses = AvgMeter()
num_batches = int(len(dataset) / config.batch_size)
tq = tqdm(range(num_batches), desc='Network Training')
for i in tq:
random_indexes = np.random.randint(len(dataset), size=config.batch_size)
input_boards, predicted_pis, predicted_vs = list(zip(*[dataset[i] for i in random_indexes]))
# Update network variables in the created placeholders
placeholders_dict = {
self.input_boards: input_boards,
self.predicted_pis: predicted_pis,
self.predicted_vs: predicted_vs,
self.dropout: config.dropout,
self.isTraining: True
}
# Start training
self.session.run(self.train_step, feed_dict=placeholders_dict)
# Calculate losses
pi_loss, v_loss = self.session.run([self.loss_pi, self.loss_v], feed_dict=placeholders_dict)
pi_losses.update(pi_loss, len(input_boards))
v_losses.update(v_loss, len(input_boards))
tq.set_postfix(Loss_pi=pi_losses, Loss_v=v_losses)
self.monitor_metrics(epoch, pi_losses.avg, v_losses.avg)
def train_one_epoch(
train_loader,
encoder,
decoder,
criterion,
encoder_optimizer,
decoder_optimizer,
device,
):
loss_meter = AvgMeter()
# switch to train mode
encoder.train()
decoder.train()
tqdm_object = tqdm(train_loader, total=len(train_loader))
for batch in tqdm_object:
batch = {k: v.to(device) for k, v in batch.items()}
batch_size = batch["images"].size(0)
features = encoder(batch["images"])
(
predictions,
encoded_captions_sorted,
decode_lengths,
alphas,
sort_ind,
) = decoder(features, batch["encoded_captions"],
batch["caption_lengths"].unsqueeze(1))
targets = encoded_captions_sorted[:, 1:]
predictions = pack_padded_sequence(predictions,
decode_lengths,
batch_first=True).data
targets = pack_padded_sequence(targets,
decode_lengths,
batch_first=True).data
loss = criterion(predictions, targets)
# record loss
loss_meter.update(loss.item(), batch_size)
loss.backward()
_ = torch.nn.utils.clip_grad_norm_(encoder.parameters(),
config.MAX_GRAD_NORM)
_ = torch.nn.utils.clip_grad_norm_(decoder.parameters(),
config.MAX_GRAD_NORM)
encoder_optimizer.step()
decoder_optimizer.step()
encoder_optimizer.zero_grad()
decoder_optimizer.zero_grad()
tqdm_object.set_postfix(train_loss=loss_meter.avg)
return loss_meter
def train(model, loader, optimizer):
tqdm_object = tqdm(loader, total=len(loader))
loss_meter = AvgMeter()
for batch in tqdm_object:
batch = {k: v.to(config.DEVICE) for k, v in batch.items()}
optimizer.zero_grad()
preds, loss = model(images=batch['images'], targets=batch['targets'])
loss.backward()
optimizer.step()
loss_meter.update(loss.item(), count=batch['images'].size(0))
tqdm_object.set_postfix(loss=loss_meter.avg)
return loss_meter
def test_epoch(net, loader):
# we transfer the mode of network to test
net.eval()
test_dice_meter = AvgMeter()
for batch_idx, (data, label) in enumerate(loader):
data = Variable(data.cuda())
output = net(data)
output = output.squeeze().data.cpu().numpy()
label = label.squeeze().cpu().numpy()
test_dice_meter.update(dice_coeff(output, label))
print("Test {} || Dice: {:.4f}".format(str(batch_idx).zfill(4), test_dice_meter.val))
return test_dice_meter.avg
def evalGivenEps(net, batch_generator, eps, nb_iter):
defense_accs = AvgMeter()
net.eval()
attack = IPGD(eps, eps / 2.0, nb_iter)
pbar = tqdm(batch_generator)
for data, label in pbar:
data = data.cuda()
label = label.cuda()
defense_accs.update(attack.get_batch_accuracy(net, data, label))
pbar.set_description('Evulating Roboustness')
return defense_accs.mean
def valid_epoch(model, valid_loader):
loss_meter = AvgMeter()
tqdm_object = tqdm(valid_loader, total=len(valid_loader))
for batch in tqdm_object:
batch = {
k: v.to(CFG.device)
for k, v in batch.items() if k != "caption"
}
loss = model(batch)
count = batch["image"].size(0)
loss_meter.update(loss.item(), count)
tqdm_object.set_postfix(valid_loss=loss_meter.avg)
return loss_meter
def val(net, loader, cost):
net.eval()
labels, predicts = [], []
loss_meter = AvgMeter()
for batch_idx, (data, label) in tqdm(enumerate(loader)): # 遍历
data = data.cuda()
label = label.cuda()
y = net(data)
loss = cost(y, label)
loss_meter.update(loss.item())
predict = y.data.cpu().numpy()
label = label.data.cpu().numpy()
predicts.extend(np.argmax(predict, axis=1))
labels.extend(label)
acc = accuracy_score(labels, predicts)
return loss_meter.avg, acc
def test_con(TaskLoader, M, SVGD, DEVICE, num_of_step=3, step_size=1e-3):
raw_M = M
LogP = AvgMeter()
pbar = tqdm(range(100))
for t in pbar:
pbar.set_description("Validation")
M = raw_M
Tasks = next(TaskLoader)
HistroyThetas = [[] for i in range(num_of_step)
] # used for recurrent kernel
for i in range(len(Tasks) - 1):
now_task = Tasks[i]
next_task = Tasks[i + 1]
X, Y, Xtest, Ytest, std = now_task
X = X.to(DEVICE)
Y = Y.to(DEVICE)
std = std.to(DEVICE) * 100 # * 100 to stablize
nextX, nextY, nextXtest, nextYtest, nextstd = next_task
nextXtest = nextXtest.to(DEVICE)
nextYtest = nextYtest.to(DEVICE)
nextstd = nextstd.to(DEVICE) * 100 # * 100 to stablize
SVGD.NablaLogP.update(X, Y, std)
SVGD.InitMomentumUpdaters()
for j in range(num_of_step):
HistroyThetas[j].append(M)
M = SVGD.step(HistroyThetas[j],
retain_graph=False,
step_size=step_size)
HistroyThetas[j][-1] = M
#HistroyThetas[j].append(M)
SVGD.NablaLogP.update(nextXtest, nextYtest, nextstd)
if i == (len(Tasks) - 2):
logp = 0
for paramsvec in M:
logp = logp + SVGD.NablaLogP(
True, paramsvec, ret_grad=False)
logp = logp / len(M)
LogP.update(logp.item())
pbar.set_postfix({'logp': LogP.mean})
return LogP.mean
def train(network, dataset, loss_fn, optimizer, epoch, writer,
early_stop=None, batch_multiplier=1):
def optimization_step(i, img, mask, lbl, save=False):
if i % batch_multiplier == 0:
optimizer.zero_grad()
if torch.cuda.is_available():
img = img.cuda()
mask = mask.cuda()
lbl = lbl.cuda()
prediction = network(img)
loss = loss_fn(prediction, mask, lbl)
loss.backward()
if (i+1) % batch_multiplier == 0:
optimizer.step()
if save:
pred = torch.sigmoid(prediction)
data = cut_to_match(pred, img)
gt = cut_to_match(pred, lbl)
writer.add_image('Train/prediction', pred[0], epoch)
writer.add_image('Train/image', data[0, :3], epoch)
writer.add_image('Train/ground_truth', gt[0], epoch)
return loss
loss_meter = AvgMeter()
network.train()
progress = tqdm(total=len(dataset), dynamic_ncols=True)
with progress:
for i, args in enumerate(dataset):
if i == early_stop:
break
loss = optimization_step(i, *args, save=i == 0)
writer.add_scalar('Train/loss', loss.item(), epoch)
progress.update(1)
loss_meter.update(loss.item())
progress.set_postfix(loss=loss_meter.last, mean=loss_meter.avg)
writer.add_scalar('Train/loss_mean', loss_meter.avg, epoch)
return loss_meter.avg
def val_epoch(net, loader, cost):
net.eval()
labels, predicts = [], []
loss_meter = AvgMeter()
for batch_idx, (feature, label) in enumerate(loader):
feature = feature.float().cuda() #GPU
label = label.float().cuda() #GPU
# feature = feature.float() #CPU
# label = label.float() #CPU
predict = net(feature)
loss = cost(predict, label.long())
loss_meter.update(loss.item())
predict = predict.data.cpu().numpy()
label = label.data.cpu().numpy()
predicts.extend(list(np.argmax(predict, axis=1)))
labels.extend(list(label))
acc = accuracy_score(labels, predicts)
return loss_meter.avg, acc
def one_epoch(
model,
criterion,
loader,
device,
optimizer=None,
lr_scheduler=None,
mode="train",
step="batch",
):
loss_meter = AvgMeter()
tqdm_object = tqdm(loader, total=len(loader))
for batch in tqdm_object:
batch = {k: v.to(device) for k, v in batch.items()}
preds = model(batch) # shape: (N, T, d_model)
caption_lengths, sort_indices = batch['caption_lengths'].sort(
dim=0, descending=True)
caption_lengths = (caption_lengths - 1).tolist()
targets = batch["encoded_captions"][sort_indices, 1:]
targets = pack_padded_sequence(targets,
caption_lengths,
batch_first=True).data
preds = pack_padded_sequence(preds, caption_lengths,
batch_first=True).data
# vocab_size = preds.size(-1)
# loss = criterion(preds.reshape(-1, vocab_size), targets.reshape(-1))
loss = criterion(preds, targets)
if mode == "train":
optimizer.zero_grad()
loss.backward()
optimizer.step()
if step == "batch":
lr_scheduler.step()
count = batch["images"].size(0)
loss_meter.update(loss.item(), count)
if mode == "train":
tqdm_object.set_postfix(loss=loss_meter.avg, lr=get_lr(optimizer))
else:
tqdm_object.set_postfix(loss=loss_meter.avg)
return loss_meter
def train_epoch(net, loader, optimizer, cost):
# we transfer the mode of network to train
net.train()
batch_loss = AvgMeter()
for batch_idx, (data, label) in enumerate(loader):
data = Variable(data.cuda())
label = Variable(label.cuda())
output = net(data)
loss = cost(output, label)
optimizer.zero_grad()
loss.backward()
optimizer.step()
batch_loss.update(loss.item())
print("Train Batch {} || Loss: {:.4f}".format(str(batch_idx).zfill(4), batch_loss.val))
return batch_loss.avg
def train_epoch(model, train_loader, optimizer, lr_scheduler, step):
loss_meter = AvgMeter()
tqdm_object = tqdm(train_loader, total=len(train_loader))
for batch in tqdm_object:
batch = {
k: v.to(CFG.device)
for k, v in batch.items() if k != "caption"
}
loss = model(batch)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if step == "batch":
lr_scheduler.step()
count = batch["image"].size(0)
loss_meter.update(loss.item(), count)
tqdm_object.set_postfix(train_loss=loss_meter.avg,
lr=get_lr(optimizer))
return loss_meter
def train(net, loader, optimizer, cost):
net.train()
loss_meter = AvgMeter()
labels, predicts = [], []
for batch_idx, (data, label) in tqdm(enumerate(loader)): # 遍历
data = data.cuda()
label = label.cuda()
optimizer.zero_grad()
y = net(data)
loss = cost(y, label)
loss_meter.update(loss.item())
loss.backward()
optimizer.step()
# 计算acc
predict = y.data.cpu().numpy()
label = label.data.cpu().numpy()
predicts.extend(np.argmax(predict, axis=1))
labels.extend(label)
if batch_idx % 10 == 0:
info = [batch_idx, loss_meter.val]
print("\rBatch: {} Loss: {:.4f}".format(*info), end="")
acc = accuracy_score(labels, predicts)
return loss_meter.avg, acc
def test(network, dataset, loss_fn, criteria, epoch, writer, early_stop=None):
network.eval()
loss_meter = AvgMeter()
progress = tqdm(total=len(dataset), dynamic_ncols=True)
with progress, torch.no_grad():
for i, (img, mask, lbl) in enumerate(dataset):
if i == early_stop:
break
if torch.cuda.is_available():
img = img.cuda()
mask = mask.cuda()
lbl = lbl.cuda()
prediction = network(img)
loss = loss_fn(prediction, mask, lbl).item()
if i == 0:
pred = torch.sigmoid(prediction)
img = cut_to_match(pred, img)
lbl = cut_to_match(pred, lbl)
writer.add_image('Test/prediction', pred[0], epoch)
writer.add_image('Test/image', img[0, :3], epoch)
writer.add_image('Test/ground_truth', lbl[0], epoch)
loss_meter.update(loss)
writer.add_scalar('Test/loss', loss, epoch)
for criterion in criteria:
value = criterion(prediction, mask, lbl)
value = value.item() if isinstance(value, torch.Tensor) else value
writer.add_scalar(f'Test/{criterion.name}', value, epoch)
progress.update(1)
writer.add_scalar('Test/loss_mean', loss_meter.avg, epoch)
def eval_one_epoch(net,
batch_generator,
device=torch.device('cuda:0'),
val_attack=None,
logger=None):
# logger.info('test start')
net.eval()
pbar = tqdm(batch_generator)
clean_accuracy = AvgMeter()
adv_accuracy = AvgMeter()
pbar.set_description('Evaluating')
for (data, label) in pbar:
data = data.to(device)
label = label.to(device)
with torch.no_grad():
pred = net(data)
acc = torch_accuracy(pred, label, (1, ))
clean_accuracy.update(acc[0].item())
if val_attack is not None:
adv_inp = val_attack.forward(data, label)
with torch.no_grad():
pred = net(adv_inp)
acc = torch_accuracy(pred, label, (1, ))
adv_accuracy.update(acc[0].item())
pbar_dic = OrderedDict()
pbar_dic['CleanAcc'] = '{:.2f}'.format(clean_accuracy.mean)
pbar_dic['AdvAcc'] = '{:.2f}'.format(adv_accuracy.mean)
pbar.set_postfix(pbar_dic)
adv_acc = adv_accuracy.mean if val_attack is not None else 0
if logger is None:
pass
else:
logger.info(
f'standard acc: {clean_accuracy.mean:.3f}%, robustness acc: {adv_accuracy.mean:.3f}%'
)
return clean_accuracy.mean, adv_acc
def train(self, dataset):
pi_losses = AvgMeter()
v_losses = AvgMeter()
for epoch in range(config.epochs):
print("Epoch: " + str(epoch))
num_batches = int(len(dataset) / config.batch_size)
for i in tqdm(range(num_batches)):
random_indexes = np.random.randint(len(dataset),
size=config.batch_size)
input_boards, target_pis, target_vs = list(
zip(*[dataset[i] for i in random_indexes]))
input_boards = np.asarray(input_boards, dtype=float)
target_pis = np.asarray(target_pis, dtype=float)
target_vs = np.asarray(target_vs, dtype=float)
losses = self.model.fit(x=input_boards,
y=[target_pis, target_vs],
batch_size=config.batch_size)
pi_losses.update(losses.history['pi_loss'], len(input_boards))
v_losses.update(losses.history['v_loss'], len(input_boards))
def train(i_epoch, network, criterion, optimizer, dataloader, device):
network.eval()
losses = []
accs = []
for idx in range(n_layers):
losses.append(AvgMeter())
accs.append(AccuracyMeter())
pbar = tqdm(dataloader)
for data in pbar:
img = data[0].to(device)
rot = data[1].long().to(device)
outputs = network(img)
for idx in range(n_layers):
outputs[idx].to(device)
optimizer.zero_grad()
all_loss = []
for idx in range(n_layers):
all_loss.append(criterion(outputs[idx], rot))
accuracy(outputs[idx], rot, accs[idx])
loss = 0
for idx in range(n_layers):
loss += all_loss[idx]
#all_loss[idx].backward()
losses[idx].add(all_loss[idx].item())
loss.backward()
optimizer.step()
lr = optimizer.param_groups[0]['lr']
str_content = generate_lossacc(n_layers, start=13)
# str_content = 'c1:{:.4f}/{:.4f} c2:{:.4f}/{:.4f} c3:{:.4f}/{:.4f} c4:{:.4f}/{:.4f} c5:{:.4f}/{:.4f}, lr:{}'
flt_content = []
for idx in range(13, n_layers):
flt_content.append(losses[idx].get())
flt_content.append(accs[idx].get())
flt_content.append(lr)
pbar.set_description("Epoch:{}".format(i_epoch))
pbar.set_postfix(info=str_content.format(*flt_content))
return losses, accs
def train(network,
dataset,
loss_fn,
optimizer,
epoch,
writer,
early_stop=None):
loss_meter = AvgMeter()
network.train()
msg = 'Train epoch {}'.format(epoch)
with tqdm(total=len(dataset), dynamic_ncols=True) as progress:
for i, args in enumerate(dataset):
if i == early_stop:
break
if torch.cuda.is_available():
args = [a.cuda() for a in args]
def evaluation():
network.zero_grad()
prediction = network(args[0])
loss = loss_fn(prediction, *args[1:])
loss.backward()
loss_meter.update(loss.item())
writer.add_scalar('Train/loss', loss.item(), epoch)
return loss
optimizer.step(evaluation)
progress.update(1)
progress.set_postfix(loss=loss_meter.last, mean=loss_meter.avg)
writer.add_scalar('Train/loss_mean', loss_meter.avg, epoch)
return loss_meter.avg
def run_ocr_48px_ctc(img: np.ndarray,
cuda: bool,
quadrilaterals: List[Tuple[Union[Quadrilateral,
TextBlock], str]],
max_chunk_size=16,
verbose: bool = False):
text_height = 48
regions = [
q.get_transformed_region(img, d, text_height)
for q, d in quadrilaterals
]
out_regions = []
perm = range(len(regions))
if len(quadrilaterals) > 0:
if isinstance(quadrilaterals[0][0], Quadrilateral):
perm = sorted(range(len(regions)),
key=lambda x: regions[x].shape[1])
ix = 0
for indices in chunks(perm, max_chunk_size):
N = len(indices)
widths = [regions[i].shape[1] for i in indices]
max_width = (4 * (max(widths) + 7) // 4) + 128
region = np.zeros((N, text_height, max_width, 3), dtype=np.uint8)
for i, idx in enumerate(indices):
W = regions[idx].shape[1]
region[i, :, :W, :] = regions[idx]
if verbose:
if quadrilaterals[idx][1] == 'v':
cv2.imwrite(
f'ocrs/{ix}.png',
cv2.rotate(
cv2.cvtColor(region[i, :, :, :],
cv2.COLOR_RGB2BGR),
cv2.ROTATE_90_CLOCKWISE))
else:
cv2.imwrite(
f'ocrs/{ix}.png',
cv2.cvtColor(region[i, :, :, :], cv2.COLOR_RGB2BGR))
ix += 1
images = (torch.from_numpy(region).float() - 127.5) / 127.5
images = einops.rearrange(images, 'N H W C -> N C H W')
if cuda:
images = images.cuda()
with torch.inference_mode():
texts = MODEL_48PX_CTC.decode(images, widths, 0, verbose=verbose)
for i, single_line in enumerate(texts):
if not single_line:
continue
cur_texts = []
total_fr = AvgMeter()
total_fg = AvgMeter()
total_fb = AvgMeter()
total_br = AvgMeter()
total_bg = AvgMeter()
total_bb = AvgMeter()
total_logprob = AvgMeter()
for (chid, logprob, fr, fg, fb, br, bg, bb) in single_line:
ch = MODEL_48PX_CTC.dictionary[chid]
if ch == '<SP>':
ch = ' '
cur_texts.append(ch)
total_logprob(logprob)
if ch != ' ':
total_fr(int(fr * 255))
total_fg(int(fg * 255))
total_fb(int(fb * 255))
total_br(int(br * 255))
total_bg(int(bg * 255))
total_bb(int(bb * 255))
prob = np.exp(total_logprob())
if prob < 0.3:
continue
txt = ''.join(cur_texts)
fr = int(total_fr())
fg = int(total_fg())
fb = int(total_fb())
br = int(total_br())
bg = int(total_bg())
bb = int(total_bb())
print(prob, txt, f'fg: ({fr}, {fg}, {fb})',
f'bg: ({br}, {bg}, {bb})')
cur_region = quadrilaterals[indices[i]][0]
if isinstance(cur_region, Quadrilateral):
cur_region.text = txt
cur_region.prob = prob
cur_region.fg_r = fr
cur_region.fg_g = fg
cur_region.fg_b = fb
cur_region.bg_r = br
cur_region.bg_g = bg
cur_region.bg_b = bb
else:
cur_region.text.append(txt)
cur_region.fg_r += fr
cur_region.fg_g += fg
cur_region.fg_b += fb
cur_region.bg_r += br
cur_region.bg_g += bg
cur_region.bg_b += bb
out_regions.append(cur_region)
return out_regions
def visualize_gt(var_map):
for kk in range(var_map.shape[0]):
pred_edge_kk = var_map[kk, :, :, :]
pred_edge_kk = pred_edge_kk.detach().cpu().numpy().squeeze()
pred_edge_kk *= 255.0
pred_edge_kk = pred_edge_kk.astype(np.uint8)
save_path = './temp/'
name = '{:02d}_gt.png'.format(kk)
cv2.imwrite(save_path + name, pred_edge_kk)
for epoch in range(1, (opt.epoch + 1)):
# scheduler.step()
generator.train()
loss_record = AvgMeter()
print('Generator Learning Rate: {}'.format(
generator_optimizer.param_groups[0]['lr']))
for i, pack in enumerate(train_loader, start=1):
for rate in size_rates:
generator_optimizer.zero_grad()
images, gts = pack
images = Variable(images)
gts = Variable(gts)
images = images.cuda()
gts = gts.cuda()
# multi-scale training samples
trainsize = int(round(opt.trainsize * rate / 32) * 32)
if rate != 1:
images = F.upsample(images,
size=(trainsize, trainsize),
semi_start_adv = 0
semi_start = 1
############################
print("Let's go!")
for epoch in range(opt.epoch):
# adjustG_lr(generator_optimizer, epoch)
# adjustD_lr(discriminator_optimizer, epoch)
adjust_lr(generator_optimizer, opt.lr_gen, epoch, opt.decay_rate,
opt.decay_epoch)
adjust_lr(discriminator_optimizer, opt.lr_dis, epoch, opt.decay_rate,
opt.decay_epoch)
#maskt = adjust_maskt(0.6, 0.2, epoch, opt.epoch)
# print('Generator Learning Rate: {}'.format(generator_optimizer.param_groups[0]['lr']))
# print('Discriminator Learning Rate: {}'.format(discriminator_optimizer.param_groups[0]['lr']))
loss_record = AvgMeter()
for i, labeled_pack in enumerate(labeled_train_loader, start=1):
generator_optimizer.zero_grad()
discriminator_optimizer.zero_grad()
generator.train()
discriminator.train()
loss_seg_value = 0
loss_adv_pred_value = 0
loss_D_value = 0
loss_semi_value = 0
loss_semi_adv_value = 0
loss_edge_value = 0
for param in discriminator.parameters():
param.requires_grad = False
fname = path + 'p{}_{}.png'.format(pred.item(), save_nr)
imageio.imsave(fname, (unnormalize(img) * 255).astype(np.uint8))
save_nr += 1
print('tracing...')
example = next(iter(train_loader))[0]
if cuda:
example = example.cuda()
net = torch.jit.trace(net, example)
print('done')
for epoch in range(n_epochs):
with tqdm(total=len(train_loader), dynamic_ncols=True) as progress:
mean_loss = AvgMeter()
mean_acc = AvgMeter()
for i, (x, y) in enumerate(train_loader):
if cuda:
x, y = x.cuda(), y.cuda()
optim.zero_grad()
maps = net(x)
predictions = maps.sum(dim=(2, 3))
loss = loss_fn(predictions, y)
loss.backward()
acc = accuracy(predictions, y)
optim.step()
progress.update(1)
mean_loss.update(loss.item())
def train_step(
model,
loader,
optimizer,
criterion,
evaluator,
vocab_size,
epoch,
teacher_forcing_ratio,
):
avg_meter = AvgMeter()
pbar = tqdm(enumerate(loader),
total=len(loader),
desc=f'Epoch {epoch} train')
now = time.time()
model.train()
for i, b in pbar:
b.src = b.src.to(device, non_blocking=non_blocking)
b.tgt = b.tgt.to(device, non_blocking=non_blocking)
b.src_mask = b.src_mask.to(device, non_blocking=non_blocking)
b.tgt_mask = b.tgt_mask.to(device, non_blocking=non_blocking)
b.gold = b.gold.to(device, non_blocking=non_blocking)
use_teacher_forcing = True if random.random(
) < teacher_forcing_ratio else False
if use_teacher_forcing:
output = model(b.src, b.tgt, b.src_mask, b.tgt_mask)
ys = output.cpu().argmax(dim=-1)
else:
ys, output = greedy_decode(model, b.src, b.src_mask, b.tgt,
b.tgt_mask, vocab_size)
scores = evaluator.get_scores(
generated_to_string(ys.numpy(), id2token), b.headline)
loss = criterion(output.view(-1, vocab_size), b.gold.view(-1))
if torch.isnan(loss):
stop = True
return
loss.backward()
optimizer.step()
model.zero_grad()
avg_meter.update('loss', loss.item())
for key, values in scores.items():
avg_meter.update(key, values)
avg_meter.inc_count()
avgs = avg_meter.get_avgs()
avg_loss = avgs['loss']
rlf = scores['rouge-l']['f']
r1f = scores['rouge-1']['f']
r2f = scores['rouge-2']['f']
avg_rlf = avgs['rouge-l']['f']
avg_r1f = avgs['rouge-1']['f']
avg_r2f = avgs['rouge-2']['f']
pbar.set_postfix(
loss=f'{loss.item():.3f} ({avg_loss:.3f})',
# rouge_l=f'{rlf:.3f} ({avg_rlf:.3f})',
rouge_1=f'{r1f:.3f} ({avg_r1f:.3f})',
rouge_2=f'{r2f:.3f} ({avg_r2f:.3f})',
)
pbar.update()
del output, loss, ys, b
if i % 100 == 0:
torch.cuda.empty_cache()
return time.time() - now, avg_meter.get_avgs()
