def train_step(self, iter, batch):
start_ind, start_pid = 0, 0
for idx, sub_batch in enumerate(batch):
if idx in self.cfg.TRAIN.unsup_dataset_indexes:
sub_batch["ind"] += start_ind
start_ind += len(self.train_sets[idx])
else:
sub_batch["ind"] = sub_batch["id"] + start_ind
start_ind += self.train_sets[idx].num_pids
sub_batch["id"] += start_pid
start_pid += self.train_sets[idx].num_pids
data = batch_processor(batch, self.cfg.MODEL.dsbn)
inputs = data["img"][0].cuda()
targets = data["id"].cuda()
indexes = data["ind"].cuda()
results = self.model(inputs)
total_loss = 0
meters = {}
for key in self.criterions.keys():
if key == "hybrid_memory":
loss = self.criterions[key](results, indexes)
else:
loss = self.criterions[key](results, targets)
total_loss += loss * float(self.cfg.TRAIN.LOSS.losses[key])
meters[key] = loss.item()
self.train_progress.update(meters)
return total_loss
def train_step(self, iter, batch):
data = batch_processor(batch, self.cfg.MODEL.dsbn)
inputs_1 = data["img"][0].cuda()
inputs_2 = data["img"][1].cuda()
targets = data["id"].cuda()
results_1, results_1_mean = self.model[0](inputs_1)
results_2, results_2_mean = self.model[1](inputs_2)
results_1["prob"] = results_1["prob"][
:, : self.train_loader.loader.dataset.num_pids
]
results_1_mean["prob"] = results_1_mean["prob"][
:, : self.train_loader.loader.dataset.num_pids
]
results_2["prob"] = results_2["prob"][
:, : self.train_loader.loader.dataset.num_pids
]
results_2_mean["prob"] = results_2_mean["prob"][
:, : self.train_loader.loader.dataset.num_pids
]
total_loss = 0
meters = {}
for key in self.criterions.keys():
if key == "soft_entropy":
loss = self.criterions[key](
results_1, results_2_mean
) + self.criterions[key](results_2, results_1_mean)
elif key == "soft_softmax_triplet":
loss = self.criterions[key](
results_1, targets, results_2_mean
) + self.criterions[key](results_2, targets, results_1_mean)
else:
loss = self.criterions[key](results_1, targets) + self.criterions[key](
results_2, targets
)
total_loss += loss * float(self.cfg.TRAIN.LOSS.losses[key])
meters[key] = loss.item()
acc_1 = accuracy(results_1["prob"].data, targets.data)
acc_2 = accuracy(results_2["prob"].data, targets.data)
meters["Acc@1"] = (acc_1[0] + acc_2[0]) * 0.5
self.train_progress.update(meters)
return total_loss