File size: 3,936 Bytes
70d1188 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 |
import os
import copy
from pathlib import Path
import torch
import torch.distributed as dist
import numpy as np
import math
import socket
# source: https://github.com/LTH14/mar/blob/main/util/misc.py
def prep_torch():
cpu_cores = get_cpu_cores()
torch.set_num_threads(cpu_cores) # intra-op threads (e.g., matrix ops)
torch.set_num_interop_threads(cpu_cores) # inter-op parallelism
os.environ["OMP_NUM_THREADS"] = str(cpu_cores)
os.environ["MKL_NUM_THREADS"] = str(cpu_cores)
os.environ["OPENBLAS_NUM_THREADS"] = str(cpu_cores)
def get_cpu_cores():
hostname = socket.gethostname()
if "bridges2" in hostname:
return int(os.environ["SLURM_JOB_CPUS_PER_NODE"])
else:
try:
with open("/sys/fs/cgroup/cpu/cpu.cfs_quota_us", "r") as f:
quota = int(f.read().strip())
with open("/sys/fs/cgroup/cpu/cpu.cfs_period_us", "r") as f:
period = int(f.read().strip())
if quota > 0:
return max(1, quota // period)
except Exception as e:
return os.cpu_count()
def setup_distributed():
dist.init_process_group(backend='nccl')
# Get the rank of the current process
rank = int(os.environ.get('RANK'))
world_size = int(os.environ.get('WORLD_SIZE'))
local_rank = int(os.environ.get('LOCAL_RANK'))
torch.cuda.set_device(local_rank)
return rank, world_size, local_rank
def is_dist_avail_and_initialized():
if not dist.is_available():
return False
if not dist.is_initialized():
return False
return True
def get_rank():
if not is_dist_avail_and_initialized():
return 0
return dist.get_rank()
def is_main_process():
return get_rank() == 0
def get_world_size():
if not is_dist_avail_and_initialized():
return 1
return dist.get_world_size()
def save_on_master(*args, **kwargs):
if is_main_process():
torch.save(*args, **kwargs)
def save_model(args, epoch, model, optimizer, ema_params=None, epoch_name=None):
if epoch_name is None:
epoch_name = str(epoch)
output_dir = Path(args.logdir)
checkpoint_path = output_dir / ('checkpoint-%s.pth' % epoch_name)
if ema_params is not None:
ema_state_dict = copy.deepcopy(model.state_dict())
for i, (name, _value) in enumerate(model.named_parameters()):
assert name in ema_state_dict
ema_state_dict[name] = ema_params[i]
else:
ema_state_dict = None
to_save = {
'model': model.state_dict(),
'optimizer': optimizer.state_dict(),
'epoch': epoch,
'args': args,
'model_ema': ema_state_dict,
}
save_on_master(to_save, checkpoint_path)
def adjust_learning_rate(optimizer, epoch, args):
"""Decay the learning rate with half-cycle cosine after warmup"""
if epoch < args.warmup_epochs:
lr = args.lr * epoch / args.warmup_epochs
else:
lr = args.min_lr + (args.lr - args.min_lr) * 0.5 * \
(1. + math.cos(math.pi * (epoch - args.warmup_epochs) / (args.n_epochs - args.warmup_epochs)))
for param_group in optimizer.param_groups:
if "lr_scale" in param_group:
param_group["lr"] = lr * param_group["lr_scale"]
else:
param_group["lr"] = lr
return lr
def add_weight_decay(model, weight_decay=1e-5, skip_list=()):
decay = []
no_decay = []
for name, param in model.named_parameters():
if not param.requires_grad:
continue # frozen weights
if len(param.shape) == 1 or name.endswith(".bias") or name in skip_list or 'diffloss' in name:
no_decay.append(param) # no weight decay on bias, norm and diffloss
else:
decay.append(param)
return [
{'params': no_decay, 'weight_decay': 0.},
{'params': decay, 'weight_decay': weight_decay}]
|