def train_step(self, imgs, labels):
imgs = imgs.cuda()
labels = [label.cuda() for label in labels]
label_sbbox, label_mbbox, label_lbbox, sbboxes, mbboxes, lbboxes = labels
outsmall, outmid, outlarge, predsmall, predmid, predlarge = self.model(
imgs)
GIOUloss, conf_loss, probloss = yololoss(self.args.MODEL, outsmall,
outmid, outlarge, predsmall,
predmid, predlarge,
label_sbbox, label_mbbox,
label_lbbox, sbboxes, mbboxes,
lbboxes)
GIOUloss = GIOUloss.sum() / imgs.shape[0]
conf_loss = conf_loss.sum() / imgs.shape[0]
probloss = probloss.sum() / imgs.shape[0]
totalloss = GIOUloss + conf_loss + probloss
self.optimizer.zero_grad()
totalloss.backward()
if self.args.Prune.sparse:
self.updateBN()
self.optimizer.step()
self.LossBox.update(GIOUloss.item())
self.LossConf.update(conf_loss.item())
self.LossClass.update(probloss.item())
def train_step_US(self, imgs, labels):
imgs = imgs.cuda()
labels = [label.cuda() for label in labels]
label_sbbox, label_mbbox, label_lbbox, sbboxes, mbboxes, lbboxes = labels
self.optimizer.zero_grad()
widths_train = []
for _ in range(4 - 2):
widths_train.append(
random.uniform(0.4, 1.0))
widths_train = [1.0,0.4] + widths_train
for idx,width_mult in enumerate(widths_train):
self.model.apply(lambda m: setattr(m, 'width_mult',width_mult))
outsmall, outmid, outlarge, predsmall, predmid, predlarge = self.model(imgs)
GIOUloss, conf_loss, probloss = yololoss(self.args.MODEL, outsmall, outmid, outlarge, predsmall, predmid,
predlarge,
label_sbbox, label_mbbox, label_lbbox, sbboxes, mbboxes, lbboxes)
GIOUloss = GIOUloss.sum() / imgs.shape[0]
conf_loss = conf_loss.sum() / imgs.shape[0]
probloss = probloss.sum() / imgs.shape[0]
totalloss = GIOUloss + conf_loss + probloss
totalloss.backward()
if idx==0:
self.LossBox.update(GIOUloss.item())
self.LossConf.update(conf_loss.item())
self.LossClass.update(probloss.item())
self.optimizer.step()
