(epoch, net, optimizer, trainloader)
| 337 | |
| 338 | |
| 339 | def train(epoch, net, optimizer, trainloader): |
| 340 | current_lr = adjust_learning_rate(optimizer, epoch) |
| 341 | train_loss = AverageMeter() |
| 342 | id_loss = AverageMeter() |
| 343 | tri_loss = AverageMeter() |
| 344 | data_time = AverageMeter() |
| 345 | batch_time = AverageMeter() |
| 346 | cnt_sum = 0 |
| 347 | id_correct = 0 |
| 348 | tri_correct = 0 |
| 349 | total = 0 |
| 350 | |
| 351 | net.train() |
| 352 | end = time.time() |
| 353 | |
| 354 | for batch_idx, (input10, input11, input2, label1, label2, true_label1, true_labels2, probV_1, probV_2, |
| 355 | probI) in enumerate(trainloader): |
| 356 | |
| 357 | labels = torch.cat((label1, label1, label2), 0) |
| 358 | input1 = torch.cat((input10, input11,), 0) |
| 359 | true_labels = torch.cat((true_label1, true_label1, true_labels2), 0) |
| 360 | prob = torch.cat((probV_1, probV_2, probI), 0) |
| 361 | input1 = input1.cuda() |
| 362 | input2 = input2.cuda() |
| 363 | labels = labels.cuda() |
| 364 | prob = prob.cuda() |
| 365 | data_time.update(time.time() - end) |
| 366 | |
| 367 | feat, out0 = net(input1, input2) |
| 368 | |
| 369 | if args.loss1 == 'sid': |
| 370 | loss_id = criterion_CE(out0, labels) |
| 371 | loss_id = prob * loss_id |
| 372 | loss_id = loss_id.sum() / prob.size(0) |
| 373 | else: |
| 374 | loss_id = criterion_id(out0, labels) |
| 375 | |
| 376 | loss_tri, batch_acc, cnt = criterion_tri(feat, out0, labels, true_labels, prob, threshold=0.5) |
| 377 | |
| 378 | loss = loss_id + loss_tri |
| 379 | |
| 380 | optimizer.zero_grad() |
| 381 | loss.backward() |
| 382 | optimizer.step() |
| 383 | |
| 384 | # update |
| 385 | train_loss.update(loss.item(), 2 * input1.size(0)) |
| 386 | id_loss.update(loss_id.item(), 2 * input1.size(0)) |
| 387 | tri_loss.update(loss_tri.item(), 2 * input1.size(0)) |
| 388 | total += labels.size(0) |
| 389 | cnt_sum += int(cnt) |
| 390 | tri_correct += batch_acc |
| 391 | _, predicted = out0.max(1) |
| 392 | id_correct += predicted.eq(labels).sum().item() |
| 393 | |
| 394 | # measure elapsed time |
| 395 | batch_time.update(time.time() - end) |
| 396 | end = time.time() |
no test coverage detected