def train(frame, label):
raw_logits, concat_logits, part_logits, _, top_n_prob = model(frame)
# Teacher Loss 来帮助计算 Navigator Loss 的临时shape
# 与partcls_loss仅仅改了一个shape
part_loss = nts_net.list_loss(
part_logits.view(batch_size * PROPOSAL_NUM, -1),
label.unsqueeze(1).repeat(1, PROPOSAL_NUM).view(-1)).view(
batch_size, PROPOSAL_NUM)
# Feature detector loss
raw_loss = criterion(raw_logits, label)
# Scrutinizing loss
concat_loss = criterion(concat_logits, label)
# Navigator loss
rank_loss = nts_net.ranking_loss(top_n_prob, part_loss)
# Teacher loss
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
raw_optimizer.zero_grad()
part_optimizer.zero_grad()
concat_optimizer.zero_grad()
partcls_optimizer.zero_grad()
total_loss.backward()
raw_optimizer.step()
part_optimizer.step()
concat_optimizer.step()
partcls_optimizer.step()
for scheduler in schedulers:
scheduler.step
scheduler.step()
# begin training
_print('--' * 50)
net.train()
for i, data in enumerate(trainloader):
niter += 1
img, label = data[0].cuda(), data[1].cuda()
batch_size = img.size(0)
raw_optimizer.zero_grad()
part_optimizer.zero_grad()
concat_optimizer.zero_grad()
partcls_optimizer.zero_grad()
raw_logits, concat_logits, part_logits, _, top_n_prob = net(img)
part_loss = model.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 = creterion(raw_logits, label)
concat_loss = creterion(concat_logits, label)
rank_loss = model.ranking_loss(top_n_prob, part_loss)
partcls_loss = creterion(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
# 总损失=整图resnet损失 + anchor attention损失 + anchor整图特征拼接损失 + anchor特征损失
total_loss.backward()
raw_optimizer.step()
part_optimizer.step()
concat_optimizer.step()
partcls_optimizer.step()
progress_bar(i, len(trainloader), 'train')