()
| 86 | |
| 87 | |
| 88 | def train(): |
| 89 | criterion = nn.CrossEntropyLoss().cuda() |
| 90 | print('train start!') |
| 91 | data_iter_s = iter(source_loader) |
| 92 | data_iter_t = iter(target_loader) |
| 93 | len_train_source = len(source_loader) |
| 94 | len_train_target = len(target_loader) |
| 95 | for step in range(conf.train.min_step + 1): |
| 96 | G.train() |
| 97 | C1.train() |
| 98 | C2.train() |
| 99 | if step % len_train_target == 0: |
| 100 | data_iter_t = iter(target_loader) |
| 101 | if step % len_train_source == 0: |
| 102 | data_iter_s = iter(source_loader) |
| 103 | data_t = next(data_iter_t) |
| 104 | data_s = next(data_iter_s) |
| 105 | inv_lr_scheduler(param_lr_g, opt_g, step, |
| 106 | init_lr=conf.train.lr, |
| 107 | max_iter=conf.train.min_step) |
| 108 | inv_lr_scheduler(param_lr_c, opt_c, step, |
| 109 | init_lr=conf.train.lr, |
| 110 | max_iter=conf.train.min_step) |
| 111 | img_s = data_s[0] |
| 112 | label_s = data_s[1] |
| 113 | img_t = data_t[0] |
| 114 | img_s, label_s = Variable(img_s.cuda()), \ |
| 115 | Variable(label_s.cuda()) |
| 116 | img_t = Variable(img_t.cuda()) |
| 117 | opt_g.zero_grad() |
| 118 | opt_c.zero_grad() |
| 119 | C2.module.weight_norm() |
| 120 | |
| 121 | ## Source loss calculation |
| 122 | feat = G(img_s) |
| 123 | out_s = C1(feat) |
| 124 | out_open = C2(feat) |
| 125 | ## source classification loss |
| 126 | loss_s = criterion(out_s, label_s) |
| 127 | ## open set loss for source |
| 128 | out_open = out_open.view(out_s.size(0), 2, -1) |
| 129 | open_loss_pos, open_loss_neg = ova_loss(out_open, label_s) |
| 130 | ## b x 2 x C |
| 131 | loss_open = 0.5 * (open_loss_pos + open_loss_neg) |
| 132 | ## open set loss for target |
| 133 | all = loss_s + loss_open |
| 134 | log_string = 'Train {}/{} \t ' \ |
| 135 | 'Loss Source: {:.4f} ' \ |
| 136 | 'Loss Open: {:.4f} ' \ |
| 137 | 'Loss Open Source Positive: {:.4f} ' \ |
| 138 | 'Loss Open Source Negative: {:.4f} ' |
| 139 | log_values = [step, conf.train.min_step, |
| 140 | loss_s.item(), loss_open.item(), |
| 141 | open_loss_pos.item(), open_loss_neg.item()] |
| 142 | if not args.no_adapt: |
| 143 | feat_t = G(img_t) |
| 144 | out_open_t = C2(feat_t) |
| 145 | out_open_t = out_open_t.view(img_t.size(0), 2, -1) |
no test coverage detected