def _train(self, epoch):
step = 0
metric = MetricTracker()
for idx, sample in enumerate(self.train_load):
step += 1
self.optimizer.zero_grad()
img = sample['img'].to(self.device)
lab = sample['lab'].float()
lab.to(self.device)
out = self.model(img)
loss = self.criterion(out, lab)
# backward
loss.backward()
self.optimizer.step()
self.lr_scheduler.step()
# updata acc&avg
metric.update_avg(loss)
metric.update_acc(out, lab)
print(f"train--step:{step}/epoch:{epoch+1}--",
f"train_loss: {metric.avg:.4f}",
f"acc:{metric.acc:.4f}",
f"lr: {self.lr_scheduler.get_lr()[0]: .2f}")
# tensorboard
self.writer.add_scalar('train_loss', metric.avg, step)
print(f'---Metrics in {epoch+1}/{self.epochs}---',
f'Training Loss : {metric.avg}',
f'Acc : {metric.acc}')
return {'loss': metric.avg, 'acc': metric.acc}
def _valid(self, epoch):
step = 0
metric = MetricTracker()
self.model.eval()
for idx, sample in enumerate(self.valie_load):
step += 1
img = sample['img'].to(self.device)
lab = sample['lab'].float()
lab.to(self.device)
out = self.model(img)
loss = self.criterion(out, lab)
# update acc&avg
metric.update_avg(loss)
metric.update_acc(out, lab)
if step % 500 == 0:
print(f"valid--step:{step}/epoch:{epoch+1}--",
f"valid_loss:{metric.avg:.4f}",
f"acc:{metric.acc:.4f}")
self.writer.add_scalar('valid_loss', metric.avg, step)
print(f'----Valid---',
f'Valid_loss:{metric.avg}',
f'Acc:{metric.acc}')
return {'valid_loss': metric.avg, 'acc': metric.acc}