def _train(self, train_loader, optimizer, epoch, total_epochs):
total_loss, correct, num_data = 0.0, 0.0, 0.0
self.model.train()
for i, data in enumerate(train_loader):
x = data["image"]
y = data["label"]
x = x.to(self.device)
y = y.to(self.device)
optimizer.zero_grad()
do_cutmix = self.cutmix and np.random.rand(1) < 0.5
if do_cutmix:
x, labels_a, labels_b, lam = cutmix_data(x=x, y=y, alpha=1.0)
logit = self.model(x)
loss = lam * self.criterion(logit, labels_a) + (
1 - lam) * self.criterion(logit, labels_b)
else:
logit = self.model(x)
loss = self.criterion(logit, y)
reg_loss = self.regularization_loss()
loss += reg_loss
loss.backward(retain_graph=True)
optimizer.step()
_, preds = logit.topk(self.topk, 1, True, True)
total_loss += loss.item()
correct += torch.sum(preds == y.unsqueeze(1)).item()
num_data += y.size(0)
n_batches = len(train_loader)
return total_loss / n_batches, correct / num_data
def update_model(self, x, y, criterion, optimizer):
optimizer.zero_grad()
do_cutmix = self.cutmix and np.random.rand(1) < 0.5
if do_cutmix:
x, labels_a, labels_b, lam = cutmix_data(x=x, y=y, alpha=1.0)
logit = self.model(x)
loss = lam * criterion(logit, labels_a) + (1 - lam) * criterion(
logit, labels_b
)
else:
logit = self.model(x)
loss = criterion(logit, y)
_, preds = logit.topk(self.topk, 1, True, True)
loss.backward()
optimizer.step()
return loss.item(), torch.sum(preds == y.unsqueeze(1)).item(), y.size(0)