| import torch
|
| import torchvision.transforms as T
|
|
|
| from tqdm import tqdm
|
| from time import sleep
|
|
|
| from utils.utils_heatmap import get_keypoints_from_heatmap_batch_maxpool
|
| from model.metrics import calculate_metrics
|
|
|
|
|
| def train_one_epoch(epoch_index, training_loader, optimizer, loss_fn, model, device):
|
| model.train(True)
|
| running_loss = 0.
|
| samples = 0
|
| transform = T.Resize((540, 960))
|
|
|
| with (tqdm(enumerate(training_loader), unit="batch", total=len(training_loader)) as tepoch):
|
| for i, data in tepoch:
|
| tepoch.set_description(f"Epoch {epoch_index}")
|
| input, target, mask = data[0].to(device), data[1].to(device), data[2].to(device)
|
| input = input if input.size()[-1] == 960 else transform(input)
|
|
|
| optimizer.zero_grad()
|
| outputs = model(input)
|
| loss = loss_fn(outputs, target) * mask.unsqueeze(-1).unsqueeze(-1)
|
| loss = loss.mean()
|
| loss.backward()
|
| optimizer.step()
|
|
|
|
|
| running_loss += loss.item()
|
| samples += mask.size()[0]
|
|
|
| tepoch.set_postfix(loss=running_loss / samples)
|
| sleep(0.1)
|
|
|
|
|
| avg_loss = running_loss / samples
|
|
|
| return avg_loss
|
|
|
|
|
| def validation_step(validation_loader, loss_fn, model, device):
|
| running_vloss = 0.0
|
|
|
| acc, prec, rec, f1 = 0, 0, 0, 0
|
| samples = 0
|
| transform = T.Resize((540, 960))
|
| model.eval()
|
|
|
| with torch.no_grad():
|
| for i, vdata in tqdm(enumerate(validation_loader), total=len(validation_loader)):
|
| input, target, mask = vdata[0].to(device), vdata[1].to(device), vdata[2].to(device)
|
| input = input if input.size()[-1] == 960 else transform(input)
|
|
|
| voutputs = model(input)
|
| vloss = loss_fn(voutputs, target) * mask.unsqueeze(-1).unsqueeze(-1)
|
| vloss = vloss.mean()
|
|
|
| kp_gt = get_keypoints_from_heatmap_batch_maxpool(target[:,:-1,:,:], return_scores=True, max_keypoints=1)
|
| kp_pred = get_keypoints_from_heatmap_batch_maxpool(voutputs[:,:-1,:,:], return_scores=True, max_keypoints=1)
|
| metrics = calculate_metrics(kp_gt, kp_pred, mask)
|
|
|
| running_vloss += vloss
|
| samples += mask.size()[0]
|
|
|
| acc += metrics[0]
|
| prec += metrics[1]
|
| rec += metrics[2]
|
| f1 += metrics[3]
|
|
|
| avg_vloss = running_vloss / samples
|
| avg_acc = acc / (i+1)
|
| avg_prec = prec / (i+1)
|
| avg_rec = rec / (i+1)
|
| avg_f1 = f1 / (i+1)
|
|
|
| return avg_vloss, avg_acc, avg_prec, avg_rec, avg_f1
|
|
|
