(args, dl, model, optimizer, loss_func, device)
| 339 | |
| 340 | |
| 341 | def train_on_epoch(args, dl, model, optimizer, loss_func, device): |
| 342 | model.train() |
| 343 | train_loss_epoch = [] |
| 344 | for data, pose in dl: |
| 345 | inputs = data.to(device) # (N, Ch, H, W) ~ (4,3,200,200), 7scenes [4, 3, 256, 341] wierd shape... |
| 346 | labels = pose.to(device) |
| 347 | if args.preprocess_ImgNet: |
| 348 | inputs = preprocess_data(inputs, device) |
| 349 | |
| 350 | predict = model(inputs) |
| 351 | loss = loss_func(predict, labels) |
| 352 | loss.backward() |
| 353 | optimizer.step() |
| 354 | optimizer.zero_grad() |
| 355 | train_loss_epoch.append(loss.item()) |
| 356 | train_loss_epoch_mean = np.mean(train_loss_epoch) |
| 357 | return train_loss_epoch_mean |
| 358 | |
| 359 | def train_posenet(args, train_dl, val_dl, model, epochs, optimizer, loss_func, scheduler, device, early_stopping): |
| 360 | writer = SummaryWriter() |
no test coverage detected