def train_model(self, **kwargs):
model = get_valid('model', kwargs)
model.train()
optimizer = get_valid('optimizer', kwargs)
epoch = get_valid('epoch', kwargs)
criterion = get_valid('criterion', kwargs)
train_loader = get_valid('train_loader', kwargs)
total_loss = 0
iter_num = len(train_loader)
log_step = iter_num // 100
log_step = max(1, log_step)
for idx, data in enumerate(train_loader):
img, label = data[0], data[1]
if self.gpu is not None:
img = img.cuda()
label = label.cuda()
model = model.cuda()
optimizer.zero_grad()
pred = model(image=img, type='classify')
loss = criterion['cls'](pred, label)
loss.backward()
optimizer.step()
if idx % log_step == 0 and idx != 0:
finish_percent = int(idx / iter_num * 100)
print('Train: finish {}%, loss: {}'.format(
finish_percent, loss.data.item()))
total_loss += loss.data.item()
avg_loss = total_loss / iter_num
log = generate_log(epoch=epoch, name='Train', avg_loss=avg_loss)
return log, avg_loss
def test_model(self, **kwargs):
model = get_valid('model', kwargs)
model.eval()
epoch = get_valid('epoch', kwargs)
criterion = get_valid('criterion', kwargs)
test_loader = get_valid('test_loader', kwargs)
total_loss = 0
iter_num = len(test_loader)
log_step = iter_num // 100
log_step = max(1, log_step)
total_correct = 0
total_pred = 0
with torch.no_grad():
for idx, (image, label, mask) in enumerate(test_loader):
target = self.encode(label, False)
target = torch.from_numpy(target)
target = target.long()
target_cp = self.encode(label, True)
target_cp = torch.from_numpy(target_cp)
target_cp = target_cp.long()
if self.use_gpu:
image = image.cuda()
target = target.cuda()
mask = mask.cuda()
target_cp = target_cp.cuda()
output = model(image=image, target=target, mask=mask)
pred_label = self.decode(output)
acc, correct, pred_num = self.compute_acc(pred_label, label)
total_correct += correct
total_pred += pred_num
output = output.contiguous().view(-1, self.class_num)
target_cp = target_cp[:, 1:].contiguous().view(-1)
loss = criterion(output, target_cp, ignore_index=-1)
total_loss += loss.data.item()
if idx % log_step == 0 and idx != 0:
finish_percent = int(idx / iter_num * 100)
print('Test: finish {}%, loss: {}'.format(
finish_percent, loss.data.item()))
acc = total_correct / total_pred
avg_loss = total_loss / iter_num
log = generate_log(epoch=epoch, name='Test', avg_loss=avg_loss)
return log, avg_loss, acc
def train_model(self, **kwargs):
super(EastDetectionTemplate, self).train_model(**kwargs)
model = get_valid('model', kwargs)
model = model.train()
optimizer = get_valid('optimizer', kwargs)
epoch = get_valid('epoch', kwargs)
criterion = get_valid('criterion', kwargs)
train_loader = get_valid('train_loader', kwargs)
total_loss = 0
aabb_loss = 0
theta_loss = 0
cls_loss = 0
iter_num = len(train_loader)
if iter_num == 0:
raise RuntimeError('training data num < batch num!')
for data in train_loader:
img = data['img']
score_map = data['score_map']
geo_map = data['geo_map']
training_mask = data['training_mask']
if self.gpu is not None:
img = img.cuda()
score_map = score_map.cuda()
geo_map = geo_map.cuda()
training_mask = training_mask.cuda()
f_score, f_geometry = model(image=img)
geo_map = geo_map.permute(0, 3, 1, 2).contiguous()
loss, l_aabb, l_theta, l_cls = criterion(score_map, f_score, geo_map, f_geometry, training_mask)
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_loss += loss.data.item()
aabb_loss += l_aabb.data.item()
theta_loss += l_theta.data.item()
cls_loss += l_cls.data.item()
avg_loss = total_loss / iter_num
avg_aabb_loss = aabb_loss / iter_num
avg_theta_loss = theta_loss / iter_num
avg_cls_loss = cls_loss / iter_num
log = generate_log(epoch=epoch,
name='Train',
avg_loss=avg_loss,
avg_aabb_loss=avg_aabb_loss,
avg_theta_loss=avg_theta_loss,
avg_cls_loss=avg_cls_loss)
return log, avg_loss
def train_model(self, **kwargs):
model = get_valid('model', kwargs)
model.train()
optimizer = get_valid('optimizer', kwargs)
epoch = get_valid('epoch', kwargs)
criterion = get_valid('criterion', kwargs)
train_loader = get_valid('train_loader', kwargs)
total_loss = 0
iter_num = len(train_loader)
log_step = iter_num // 100
log_step = max(1, log_step)
if iter_num == 0:
raise RuntimeError('training data num < batch num!')
for idx, (image, label, mask) in enumerate(train_loader):
target = self.encode(label, False)
target = torch.from_numpy(target)
target = target.long()
target_cp = self.encode(label, True)
target_cp = torch.from_numpy(target_cp)
target_cp = target_cp.long()
if self.use_gpu:
image = image.cuda()
mask = mask.cuda()
target = target.cuda()
target_cp = target_cp.cuda()
output = model(image=image, target=target, mask=mask)
output = output.contiguous().view(-1, self.class_num)
target_cp = target_cp[:, 1:].contiguous().view(-1)
loss = criterion(output, target_cp, ignore_index=-1)
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_loss += loss.data.item()
if idx % log_step == 0 and idx != 0:
finish_percent = int(idx / iter_num * 100)
print('Train: finish {}%, loss: {}'.format(
finish_percent, loss.data.item()))
avg_loss = total_loss / iter_num
log = generate_log(epoch=epoch, name='Train', avg_loss=avg_loss)
return log, avg_loss
def test_model(self, **kwargs):
super(EastDetectionTemplate, self).test_model(**kwargs)
model = get_valid('model', kwargs)
model = model.eval()
epoch = get_valid('epoch', kwargs)
criterion = get_valid('criterion', kwargs)
test_loader = get_valid('test_loader', kwargs)
total_loss = 0
aabb_loss = 0
theta_loss = 0
cls_loss = 0
iter_num = len(test_loader)
with torch.no_grad():
for data in test_loader:
img = data['img']
score_map = data['score_map']
geo_map = data['geo_map']
training_mask = data['training_mask']
# img_path = data['img_path']
if self.gpu is not None:
img = img.cuda()
score_map = score_map.cuda()
geo_map = geo_map.cuda()
training_mask = training_mask.cuda()
f_score, f_geometry = model(image=img)
geo_map = geo_map.permute(0, 3, 1, 2).contiguous()
loss, l_aabb, l_theta, l_cls = criterion(score_map, f_score, geo_map, f_geometry, training_mask)
total_loss += loss.data.item()
aabb_loss += l_aabb.data.item()
theta_loss += l_theta.data.item()
cls_loss += l_cls.data.item()
avg_loss = total_loss / iter_num
avg_aabb_loss = aabb_loss / iter_num
avg_theta_loss = theta_loss / iter_num
avg_cls_loss = cls_loss / iter_num
log = generate_log(epoch=epoch,
name='Test',
avg_loss=avg_loss,
avg_aabb_loss=avg_aabb_loss,
avg_theta_loss=avg_theta_loss,
avg_cls_loss=avg_cls_loss)
return log, avg_loss
def train_model(self, **kwargs):
super(PseDetectionTemplate, self).train_model(**kwargs)
model = get_valid('model', kwargs)
model.train()
optimizer = get_valid('optimizer', kwargs)
epoch = get_valid('epoch', kwargs)
criterion = get_valid('criterion', kwargs)
train_loader = get_valid('train_loader', kwargs)
total_loss = 0
total_text_loss = 0
total_kernels_loss = 0
iter_num = len(train_loader)
if iter_num == 0:
raise RuntimeError('training data num < batch num!')
for data in train_loader:
img = data['img']
score_maps = data['score_maps']
training_mask = data['training_mask']
if self.gpu is not None:
img = img.cuda()
score_maps = score_maps.cuda()
training_mask = training_mask.cuda()
output = model(image=img)
loss_text, loss_kernels, loss = criterion(output, score_maps, training_mask)
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_loss += loss.data.item()
total_text_loss += loss_text.data.item()
total_kernels_loss += loss_kernels.data.item()
avg_loss = total_loss / iter_num
avg_text_loss = total_text_loss / iter_num
avg_kernels_loss = total_kernels_loss / iter_num
log = generate_log(epoch=epoch,
name='Train',
avg_loss=avg_loss,
avg_text_loss=avg_text_loss,
avg_kernels_loss=avg_kernels_loss)
return log, avg_loss
def test_model(self, **kwargs):
super(ClassifyTemplate, self).test_model(**kwargs)
test_loader = get_valid('test_loader', kwargs)
criterion = get_valid('criterion', kwargs)
model = get_valid('model', kwargs)
epoch = get_valid('epoch', kwargs)
sum_loss = 0
correct = 0
iter_num = len(test_loader)
log_step = iter_num // 100
log_step = max(1, log_step)
model = model.eval()
with torch.no_grad():
for idx, data in enumerate(test_loader):
img, label = data[0], data[1]
if self.gpu is not None:
img = img.cuda()
label = label.cuda()
model = model.cuda()
pred = model(image=img)
loss = criterion(pred, label)
_, pred_cls = torch.max(pred, 1)
correct += torch.sum(pred_cls.data == label.data)
sum_loss += loss.data.item()
if idx % log_step == 0 and idx != 0:
finish_percent = int(idx / iter_num * 100)
print('Test: finish {}%, loss: {}'.format(
finish_percent, loss.data.item()))
avg_loss = sum_loss / iter_num
acc = int(correct) / self.test_data_num
log = generate_log(epoch=epoch,
name='Test',
avg_loss=avg_loss,
acc=acc)
return log, avg_loss, acc
def test_model(self, **kwargs):
super(PseDetectionTemplate, self).test_model(**kwargs)
model = get_valid('model', kwargs)
model.train()
epoch = get_valid('epoch', kwargs)
criterion = get_valid('criterion', kwargs)
test_loader = get_valid('test_loader', kwargs)
total_loss = 0
total_text_loss = 0
total_kernels_loss = 0
iter_num = len(test_loader)
with torch.no_grad():
for data in test_loader:
img = data['img']
score_maps = data['score_maps']
training_mask = data['training_mask']
if self.gpu is not None:
img = img.cuda()
score_maps = score_maps.cuda()
training_mask = training_mask.cuda()
output = model(image=img)
loss_text, loss_kernels, loss = criterion(output, score_maps, training_mask)
total_loss += loss.data.item()
total_text_loss += loss_text.data.item()
total_kernels_loss += loss_kernels.data.item()
avg_loss = total_loss / iter_num
avg_text_loss = total_text_loss / iter_num
avg_kernels_loss = total_kernels_loss / iter_num
log = generate_log(epoch=epoch,
name='Test',
avg_loss=avg_loss,
avg_text_loss=avg_text_loss,
avg_kernels_loss=avg_kernels_loss)
return log, avg_loss
def test_model(self, **kwargs):
test_loader = get_valid('test_loader', kwargs)
criterion = get_valid('criterion', kwargs)
model = get_valid('model', kwargs)
epoch = get_valid('epoch', kwargs)
loss = 0
correct = 0
iter_num = len(test_loader)
log_step = iter_num // 100
log_step = max(1, log_step)
model = model.eval()
with torch.no_grad():
for idx, data in enumerate(test_loader):
img, label = data[0], data[1]
if self.gpu is not None:
img = img.cuda()
label = label.cuda()
batch_size = img.size(0)
_, concat_logits, _, _, _ = model(image=img)
concat_loss = criterion(concat_logits, label)
_, concat_predict = torch.max(concat_logits, 1)
correct += torch.sum(concat_predict.data == label.data)
loss += concat_loss.item() * batch_size
if idx % log_step == 0 and idx != 0:
finish_percent = int(idx / iter_num * 100)
print('Test: finish {}%,concat loss: {},'.format(
finish_percent, float(concat_loss.item())))
loss = loss / self.test_data_num
acc = float(correct) / self.test_data_num
log = generate_log(epoch=epoch, name='Test', loss=loss, acc=acc)
return log, loss, acc
def test_model(self, **kwargs):
model = get_valid('model', kwargs)
model.eval()
epoch = get_valid('epoch', kwargs)
criterion = get_valid('criterion', kwargs)
test_loader = get_valid('test_loader', kwargs)
total_loss = 0
total_correct = 0
total_pred = 0
class_dict = get_valid('class_dict', kwargs)
for cls in class_dict:
cls_label = class_dict[cls]
data_loader = test_loader[cls_label]
iter_num = len(data_loader)
log_step = iter_num // 100
log_step = max(1, log_step)
class_correct = 0
class_pred = 0
with torch.no_grad():
for idx, (image, label, mask) in enumerate(data_loader):
target = self.encode(label, False, cls_label)
target = torch.from_numpy(target)
target = target.long()
target_cp = self.encode(label, True, cls_label)
target_cp = torch.from_numpy(target_cp)
target_cp = target_cp.long()
batch_size = image.shape[0]
cls_tensor = torch.zeros((batch_size))
cls_tensor += cls_label
cls_tensor = cls_tensor.long()
if self.use_gpu:
image = image.cuda()
target = target.cuda()
mask = mask.cuda()
target_cp = target_cp.cuda()
cls_tensor = cls_tensor.cuda()
output, cls_pred = model(image=image,
target=target,
mask=mask,
cls_label=cls_label,
type='sar')
pred_label = self.decode(output, cls_label)
acc, correct, pred_num = self.compute_acc(
pred_label, label)
total_correct += correct
total_pred += pred_num
class_correct += correct
class_pred += pred_num
class_num = self.class_num_list[cls_label]
output = output.contiguous().view(-1, class_num)
target_cp = target_cp[:, 1:].contiguous().view(-1)
sar_loss = criterion['sar'](output,
target_cp,
ignore_index=-1)
cls_loss = criterion['cls'](cls_pred, cls_tensor)
loss = sar_loss + cls_loss
total_loss += loss.data.item()
if idx % log_step == 0 and idx != 0:
finish_percent = int(idx / iter_num * 100)
print(
'Test {}: finish {}%, loss: {}, sar loss:{}, cls loss: {}'
.format(cls, finish_percent, loss.data.item(),
sar_loss.item(), cls_loss.item()))
print('Test {}: acc: {}'.format(cls,
class_correct / class_pred))
acc = total_correct / total_pred
avg_loss = total_loss / iter_num
log = generate_log(epoch=epoch, name='Test', avg_loss=avg_loss)
return log, avg_loss, acc
def train_model(self, **kwargs):
model = get_valid('model', kwargs)
model.train()
optimizer = get_valid('optimizer', kwargs)
epoch = get_valid('epoch', kwargs)
criterion = get_valid('criterion', kwargs)
train_loader = get_valid('train_loader', kwargs)
total_loss = 0
class_dict = get_valid('class_dict', kwargs)
key_list = list()
for key in class_dict:
key_list.append(key)
cls = random.randint(0, len(key_list) - 1)
cls = key_list[cls]
cls_label = class_dict[cls]
data_loader = train_loader[cls_label]['data_loader']
repeat_num = train_loader[cls_label]['repeat_num']
iter_num = len(data_loader)
log_step = iter_num // 100
log_step = max(1, log_step)
model.fc.weight.requires_grad = False
model.encoder.requires_grad = False
if iter_num == 0:
raise RuntimeError('training data num < batch num!')
for idy in range(repeat_num):
for idx, (image, label, mask) in enumerate(data_loader):
target = self.encode(label, False, cls_label)
target = torch.from_numpy(target)
target = target.long()
target_cp = self.encode(label, True, cls_label)
target_cp = torch.from_numpy(target_cp)
target_cp = target_cp.long()
batch_size = image.shape[0]
cls_tensor = torch.zeros((batch_size))
cls_tensor += cls_label
cls_tensor = cls_tensor.long()
if self.use_gpu:
image = image.cuda()
mask = mask.cuda()
target = target.cuda()
target_cp = target_cp.cuda()
cls_tensor = cls_tensor.cuda()
output, cls_pred = model(image=image,
target=target,
mask=mask,
cls_label=cls_label,
type='sar')
class_num = self.class_num_list[cls_label]
output = output.contiguous().view(-1, class_num)
target_cp = target_cp[:, 1:].contiguous().view(-1)
sar_loss = criterion['sar'](output, target_cp, ignore_index=-1)
cls_loss = criterion['cls'](cls_pred, cls_tensor)
loss = sar_loss + cls_loss
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_loss += loss.data.item()
if idx % log_step == 0 and idx != 0:
finish_percent = int(idx / iter_num * 100)
print(
'Train {}: {}/{} finish {}%, loss: {}, sar loss:{}, cls loss: {}'
.format(cls, idy, repeat_num, finish_percent,
loss.data.item(), sar_loss.item(),
cls_loss.item()))
avg_loss = total_loss / iter_num
log = generate_log(epoch=epoch, name='Train', avg_loss=avg_loss)
return log, avg_loss
def train_model(self, **kwargs):
(raw_optimizer, concat_optimizer, part_optimizer, partcls_optimizer,
schedulers) = get_valid('optimizer', kwargs)
train_loader = get_valid('train_loader', kwargs)
criterion = get_valid('criterion', kwargs)
proposal_num = get('proposal_num', kwargs, 6)
model = get_valid('model', kwargs)
epoch = get_valid('epoch', kwargs)
sum_total_loss = 0
sum_concat_loss = 0
sum_raw_loss = 0
sum_rank_loss = 0
sum_partcls_loss = 0
iter_num = len(train_loader)
log_step = iter_num // 100
log_step = max(1, log_step)
for scheduler in schedulers:
scheduler.step()
model = model.train()
for idx, data in enumerate(train_loader):
img, label = data[0], data[1]
if self.gpu is not None:
img = img.cuda()
label = label.cuda()
raw_optimizer.zero_grad()
concat_optimizer.zero_grad()
part_optimizer.zero_grad()
partcls_optimizer.zero_grad()
batch_size = img.size(0)
raw_logits, concat_logits, part_logits, _, top_n_prob = model(
image=img)
part_loss = NTSClassifyModel.list_loss(
part_logits.view(batch_size * proposal_num, -1),
label.unsqueeze(1).repeat(1, proposal_num).view(-1)).view(
batch_size, proposal_num)
raw_loss = criterion(raw_logits, label)
concat_loss = criterion(concat_logits, label)
rank_loss = NTSClassifyModel.ranking_loss(top_n_prob, part_loss,
proposal_num, **kwargs)
partcls_loss = criterion(
part_logits.view(batch_size * proposal_num, -1),
label.unsqueeze(1).repeat(1, proposal_num).view(-1))
total_loss = raw_loss + rank_loss + concat_loss + partcls_loss
total_loss.backward()
raw_optimizer.step()
part_optimizer.step()
concat_optimizer.step()
partcls_optimizer.step()
sum_total_loss += total_loss.data.item()
sum_raw_loss += raw_loss.data.item()
sum_concat_loss += concat_loss.data.item()
sum_rank_loss += rank_loss.data.item()
sum_partcls_loss += partcls_loss.data.item()
if idx % log_step == 0 and idx != 0:
finish_percent = int(idx / iter_num * 100)
print(
'Train: finish {}%,total loss: {}, raw loss: {}, rank loss: {}, concat loss: {},'
' partcls loss: {}'.format(finish_percent,
float(total_loss),
float(raw_loss),
float(rank_loss),
float(concat_loss),
float(partcls_loss)))
avg_total_loss = sum_total_loss / iter_num
avg_raw_loss = sum_raw_loss / iter_num
avg_concat_loss = sum_concat_loss / iter_num
avg_rank_loss = sum_rank_loss / iter_num
avg_partcls_loss = sum_partcls_loss / iter_num
log = generate_log(epoch=epoch,
name='Train',
avg_total_loss=avg_total_loss,
avg_raw_loss=avg_raw_loss,
avg_concat_loss=avg_concat_loss,
avg_rank_loss=avg_rank_loss,
avg_partcls_loss=avg_partcls_loss)
return log, avg_total_loss