298 lines
8.8 KiB
Python
Executable File
298 lines
8.8 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
|
|
## Copyright 2018 Intel Corporation
|
|
## SPDX-License-Identifier: Apache-2.0
|
|
|
|
import os
|
|
import sys
|
|
from glob import glob
|
|
import time
|
|
import numpy as np
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.optim as optim
|
|
import torch.distributed as dist
|
|
from torch.utils.tensorboard import SummaryWriter
|
|
|
|
from config import *
|
|
from dataset import *
|
|
from model import *
|
|
from loss import *
|
|
from result import *
|
|
from util import *
|
|
|
|
def main():
|
|
# Parse the command line arguments
|
|
cfg = parse_args(description='Trains a model using preprocessed datasets.')
|
|
|
|
# Start the worker(s)
|
|
start_workers(cfg, main_worker)
|
|
|
|
# Worker function
|
|
def main_worker(rank, cfg):
|
|
# Initialize the worker
|
|
distributed = init_worker(rank, cfg)
|
|
|
|
# Initialize the random seed
|
|
if cfg.seed is not None:
|
|
torch.manual_seed(cfg.seed)
|
|
|
|
# Initialize the PyTorch device
|
|
device_id = cfg.device_id + rank
|
|
device = init_device(cfg, id=device_id)
|
|
|
|
# Initialize the model
|
|
model = orig_model = get_model(cfg)
|
|
model.to(device)
|
|
if cfg.compile:
|
|
model = torch.compile(model, mode=cfg.compile)
|
|
if distributed:
|
|
model = nn.parallel.DistributedDataParallel(model, device_ids=[device_id])
|
|
|
|
# Initialize the loss function
|
|
loss_fn = get_loss_function(cfg)
|
|
loss_fn.to(device)
|
|
if cfg.compile:
|
|
loss_fn = torch.compile(loss_fn, mode=cfg.compile)
|
|
|
|
# Initialize the optimizer
|
|
optimizer = optim.Adam(model.parameters(), lr=1)
|
|
|
|
# Check whether the result already exists
|
|
result_dir = get_result_dir(cfg)
|
|
resume = os.path.isdir(result_dir)
|
|
|
|
# Sync the workers (required due to the previous isdir check)
|
|
if distributed:
|
|
dist.barrier()
|
|
|
|
# Start or resume training
|
|
if resume:
|
|
if rank == 0:
|
|
print('Resuming result:', cfg.result)
|
|
|
|
# Load and verify the config
|
|
result_cfg = load_config(result_dir)
|
|
if set(result_cfg.features) != set(cfg.features):
|
|
error('input feature set mismatch')
|
|
|
|
# Restore the latest checkpoint
|
|
last_epoch = get_latest_checkpoint_epoch(result_dir)
|
|
checkpoint = load_checkpoint(result_dir, device, last_epoch, orig_model, optimizer)
|
|
step = checkpoint['step']
|
|
else:
|
|
if rank == 0:
|
|
print('Result:', cfg.result)
|
|
os.makedirs(result_dir)
|
|
|
|
# Save the config
|
|
save_config(result_dir, cfg)
|
|
|
|
# Save the source code
|
|
src_filenames = glob(os.path.join(os.path.dirname(sys.argv[0]), '*.py'))
|
|
src_zip_filename = os.path.join(result_dir, 'src.zip')
|
|
save_zip(src_zip_filename, src_filenames)
|
|
|
|
last_epoch = 0
|
|
step = 0
|
|
|
|
# Make sure all workers have loaded the checkpoint
|
|
if distributed:
|
|
dist.barrier()
|
|
|
|
start_epoch = last_epoch + 1
|
|
if start_epoch > cfg.num_epochs:
|
|
exit() # nothing to do
|
|
|
|
# Reset the random seed if resuming result
|
|
if cfg.seed is not None and start_epoch > 1:
|
|
seed = cfg.seed + start_epoch - 1
|
|
torch.manual_seed(seed)
|
|
|
|
# Initialize the training dataset
|
|
train_data = TrainingDataset(cfg, cfg.train_data)
|
|
if len(train_data) > 0:
|
|
if rank == 0:
|
|
print('Training images:', train_data.num_images)
|
|
else:
|
|
error('no training images')
|
|
train_loader, train_sampler = get_data_loader(rank, cfg, train_data, shuffle=True)
|
|
train_steps_per_epoch = len(train_loader)
|
|
|
|
# Initialize the validation dataset
|
|
valid_data = ValidationDataset(cfg, cfg.valid_data)
|
|
if len(valid_data) > 0:
|
|
if rank == 0:
|
|
print('Validation images:', valid_data.num_images)
|
|
valid_loader, valid_sampler = get_data_loader(rank, cfg, valid_data, shuffle=False)
|
|
valid_steps_per_epoch = len(valid_loader)
|
|
else:
|
|
valid_steps_per_epoch = 0
|
|
|
|
# Initialize the learning rate scheduler
|
|
lr_scheduler = optim.lr_scheduler.OneCycleLR(
|
|
optimizer,
|
|
max_lr=cfg.max_lr,
|
|
total_steps=cfg.num_epochs,
|
|
pct_start=cfg.lr_warmup,
|
|
anneal_strategy='cos',
|
|
div_factor=(25. if cfg.lr is None else cfg.max_lr / cfg.lr),
|
|
final_div_factor=1e4,
|
|
last_epoch=last_epoch-1)
|
|
|
|
if lr_scheduler.last_epoch != last_epoch:
|
|
error('failed to restore LR scheduler state')
|
|
|
|
# Check whether AMP is enabled
|
|
amp_enabled = cfg.precision == 'mixed'
|
|
|
|
if amp_enabled:
|
|
# Initialize the gradient scaler
|
|
scaler = torch.amp.GradScaler()
|
|
|
|
# Initialize the summary writer
|
|
log_dir = get_result_log_dir(result_dir)
|
|
if rank == 0:
|
|
summary_writer = SummaryWriter(log_dir)
|
|
if step == 0:
|
|
summary_writer.add_scalar('learning_rate', lr_scheduler.get_last_lr()[0], 0)
|
|
|
|
# Training and evaluation loops
|
|
if rank == 0:
|
|
print()
|
|
progress_format = '%-5s %' + str(len(str(cfg.num_epochs))) + 'd/%d:' % cfg.num_epochs
|
|
total_start_time = time.time()
|
|
|
|
for epoch in range(start_epoch, cfg.num_epochs+1):
|
|
if rank == 0:
|
|
start_time = time.time()
|
|
progress = ProgressBar(train_steps_per_epoch, progress_format % ('Train', epoch))
|
|
|
|
# Switch to training mode
|
|
model.train()
|
|
train_loss = 0.
|
|
|
|
# Iterate over the batches
|
|
if distributed:
|
|
train_sampler.set_epoch(epoch)
|
|
|
|
for i, batch in enumerate(train_loader, 0):
|
|
if cfg.device == 'cuda' and cfg.compile in {'reduce-overhead', 'max-autotune'}:
|
|
torch.compiler.cudagraph_mark_step_begin()
|
|
|
|
# Get the batch
|
|
input, target = batch
|
|
input = input.to(device, non_blocking=True)
|
|
target = target.to(device, non_blocking=True)
|
|
if not amp_enabled:
|
|
input = input.float()
|
|
target = target.float()
|
|
|
|
# Run a training step
|
|
optimizer.zero_grad()
|
|
|
|
with torch.autocast('cuda', enabled=amp_enabled):
|
|
loss = loss_fn(model(input), target)
|
|
|
|
if amp_enabled:
|
|
scaler.scale(loss).backward()
|
|
scaler.step(optimizer)
|
|
scaler.update()
|
|
else:
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
# Next step
|
|
step += 1
|
|
train_loss += loss.detach()
|
|
if rank == 0:
|
|
progress.next()
|
|
|
|
# Get and update the learning rate
|
|
lr = lr_scheduler.get_last_lr()[0]
|
|
lr_scheduler.step()
|
|
|
|
# Compute the average training loss
|
|
if distributed:
|
|
dist.all_reduce(train_loss, op=dist.ReduceOp.SUM)
|
|
train_loss = train_loss.item() / (train_steps_per_epoch * cfg.num_devices)
|
|
|
|
# Write summary
|
|
if rank == 0:
|
|
summary_writer.add_scalar('learning_rate', lr, epoch)
|
|
summary_writer.add_scalar('loss', train_loss, epoch)
|
|
|
|
# Print stats
|
|
if rank == 0:
|
|
duration = time.time() - start_time
|
|
total_duration = time.time() - total_start_time
|
|
images_per_sec = train_steps_per_epoch * cfg.batch_size / duration
|
|
eta = ((cfg.num_epochs - epoch) * total_duration / (epoch + 1 - start_epoch))
|
|
progress.finish('loss=%.6f, lr=%.6f (%.1f images/s, %s, eta %s)'
|
|
% (train_loss, lr, images_per_sec, format_time(duration), format_time(eta, precision=2)))
|
|
|
|
if ((cfg.num_valid_epochs > 0 and epoch % cfg.num_valid_epochs == 0) or epoch == cfg.num_epochs) \
|
|
and valid_steps_per_epoch > 0:
|
|
# Validation
|
|
if rank == 0:
|
|
start_time = time.time()
|
|
progress = ProgressBar(valid_steps_per_epoch, progress_format % ('Valid', epoch))
|
|
|
|
# Switch to evaluation mode
|
|
model.eval()
|
|
valid_loss = 0.
|
|
|
|
# Iterate over the batches
|
|
with torch.no_grad():
|
|
for _, batch in enumerate(valid_loader, 0):
|
|
if cfg.device == 'cuda' and cfg.compile in {'reduce-overhead', 'max-autotune'}:
|
|
torch.compiler.cudagraph_mark_step_begin()
|
|
|
|
# Get the batch
|
|
input, target = batch
|
|
input = input.to(device, non_blocking=True)
|
|
target = target.to(device, non_blocking=True)
|
|
if not amp_enabled:
|
|
input = input.float()
|
|
target = target.float()
|
|
|
|
# Run a validation step
|
|
with torch.autocast('cuda', enabled=amp_enabled):
|
|
loss = loss_fn(model(input), target)
|
|
|
|
# Next step
|
|
valid_loss += loss.detach()
|
|
if rank == 0:
|
|
progress.next()
|
|
|
|
# Compute the average validation loss
|
|
if distributed:
|
|
dist.all_reduce(valid_loss, op=dist.ReduceOp.SUM)
|
|
valid_loss = valid_loss.item() / (valid_steps_per_epoch * cfg.num_devices)
|
|
|
|
# Write summary
|
|
if rank == 0:
|
|
summary_writer.add_scalar('valid_loss', valid_loss, epoch)
|
|
|
|
# Print stats
|
|
if rank == 0:
|
|
duration = time.time() - start_time
|
|
images_per_sec = valid_steps_per_epoch * cfg.batch_size / duration
|
|
progress.finish('valid_loss=%.6f (%.1f images/s, %.1fs)'
|
|
% (valid_loss, images_per_sec, duration))
|
|
|
|
if (rank == 0) and ((cfg.num_save_epochs > 0 and epoch % cfg.num_save_epochs == 0) or epoch == cfg.num_epochs):
|
|
# Save a checkpoint
|
|
save_checkpoint(result_dir, epoch, step, orig_model, optimizer)
|
|
|
|
# Print final stats
|
|
if rank == 0:
|
|
total_duration = time.time() - total_start_time
|
|
print('\nFinished (%s)' % format_time(total_duration))
|
|
|
|
# Cleanup
|
|
cleanup_worker(cfg)
|
|
|
|
if __name__ == '__main__':
|
|
main() |