|
|
"""
|
|
|
Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
|
|
|
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
|
|
|
"""
|
|
|
|
|
|
import torch
|
|
|
import torch.nn as nn
|
|
|
|
|
|
from ..misc import MetricLogger, SmoothedValue, reduce_dict
|
|
|
|
|
|
|
|
|
def train_one_epoch(
|
|
|
model: nn.Module, criterion: nn.Module, dataloader, optimizer, ema, epoch, device
|
|
|
):
|
|
|
""" """
|
|
|
model.train()
|
|
|
|
|
|
metric_logger = MetricLogger(delimiter=" ")
|
|
|
metric_logger.add_meter("lr", SmoothedValue(window_size=1, fmt="{value:.6f}"))
|
|
|
print_freq = 100
|
|
|
header = "Epoch: [{}]".format(epoch)
|
|
|
|
|
|
for imgs, labels in metric_logger.log_every(dataloader, print_freq, header):
|
|
|
imgs = imgs.to(device)
|
|
|
labels = labels.to(device)
|
|
|
|
|
|
preds = model(imgs)
|
|
|
loss: torch.Tensor = criterion(preds, labels, epoch)
|
|
|
|
|
|
optimizer.zero_grad()
|
|
|
loss.backward()
|
|
|
optimizer.step()
|
|
|
|
|
|
if ema is not None:
|
|
|
ema.update(model)
|
|
|
|
|
|
loss_reduced_values = {k: v.item() for k, v in reduce_dict({"loss": loss}).items()}
|
|
|
metric_logger.update(**loss_reduced_values)
|
|
|
metric_logger.update(lr=optimizer.param_groups[0]["lr"])
|
|
|
|
|
|
metric_logger.synchronize_between_processes()
|
|
|
print("Averaged stats:", metric_logger)
|
|
|
|
|
|
stats = {k: meter.global_avg for k, meter in metric_logger.meters.items()}
|
|
|
return stats
|
|
|
|
|
|
|
|
|
@torch.no_grad()
|
|
|
def evaluate(model, criterion, dataloader, device):
|
|
|
model.eval()
|
|
|
|
|
|
metric_logger = MetricLogger(delimiter=" ")
|
|
|
|
|
|
|
|
|
metric_logger.add_meter("acc", SmoothedValue(window_size=1))
|
|
|
metric_logger.add_meter("loss", SmoothedValue(window_size=1))
|
|
|
|
|
|
header = "Test:"
|
|
|
for imgs, labels in metric_logger.log_every(dataloader, 10, header):
|
|
|
imgs, labels = imgs.to(device), labels.to(device)
|
|
|
preds = model(imgs)
|
|
|
|
|
|
acc = (preds.argmax(dim=-1) == labels).sum() / preds.shape[0]
|
|
|
loss = criterion(preds, labels)
|
|
|
|
|
|
dict_reduced = reduce_dict({"acc": acc, "loss": loss})
|
|
|
reduced_values = {k: v.item() for k, v in dict_reduced.items()}
|
|
|
metric_logger.update(**reduced_values)
|
|
|
|
|
|
metric_logger.synchronize_between_processes()
|
|
|
print("Averaged stats:", metric_logger)
|
|
|
|
|
|
stats = {k: meter.global_avg for k, meter in metric_logger.meters.items()}
|
|
|
return stats
|
|
|
|
