Mido Assran
157 lines
4.5 KiB
Python
157 lines
4.5 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
#
|
|
|
|
import logging
|
|
import sys
|
|
|
|
import torch
|
|
|
|
import src.models.vision_transformer as vit
|
|
from src.utils.schedulers import (
|
|
WarmupCosineSchedule,
|
|
CosineWDSchedule)
|
|
from src.utils.tensors import trunc_normal_
|
|
|
|
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
|
|
logger = logging.getLogger()
|
|
|
|
|
|
def load_checkpoint(
|
|
device,
|
|
r_path,
|
|
encoder,
|
|
predictor,
|
|
target_encoder,
|
|
opt,
|
|
scaler,
|
|
):
|
|
try:
|
|
checkpoint = torch.load(r_path, map_location=torch.device('cpu'))
|
|
epoch = checkpoint['epoch']
|
|
|
|
# -- loading encoder
|
|
pretrained_dict = checkpoint['encoder']
|
|
msg = encoder.load_state_dict(pretrained_dict)
|
|
logger.info(f'loaded pretrained encoder from epoch {epoch} with msg: {msg}')
|
|
|
|
# -- loading predictor
|
|
pretrained_dict = checkpoint['predictor']
|
|
msg = predictor.load_state_dict(pretrained_dict)
|
|
logger.info(f'loaded pretrained encoder from epoch {epoch} with msg: {msg}')
|
|
|
|
# -- loading target_encoder
|
|
if target_encoder is not None:
|
|
print(list(checkpoint.keys()))
|
|
pretrained_dict = checkpoint['target_encoder']
|
|
msg = target_encoder.load_state_dict(pretrained_dict)
|
|
logger.info(f'loaded pretrained encoder from epoch {epoch} with msg: {msg}')
|
|
|
|
# -- loading optimizer
|
|
opt.load_state_dict(checkpoint['opt'])
|
|
if scaler is not None:
|
|
scaler.load_state_dict(checkpoint['scaler'])
|
|
logger.info(f'loaded optimizers from epoch {epoch}')
|
|
logger.info(f'read-path: {r_path}')
|
|
del checkpoint
|
|
|
|
except Exception as e:
|
|
logger.info(f'Encountered exception when loading checkpoint {e}')
|
|
epoch = 0
|
|
|
|
return encoder, predictor, target_encoder, opt, scaler, epoch
|
|
|
|
|
|
def init_model(
|
|
device,
|
|
patch_size=16,
|
|
model_name='vit_base',
|
|
crop_size=224,
|
|
pred_depth=6,
|
|
pred_emb_dim=384
|
|
):
|
|
encoder = vit.__dict__[model_name](
|
|
img_size=[crop_size],
|
|
patch_size=patch_size)
|
|
predictor = vit.__dict__['vit_predictor'](
|
|
num_patches=encoder.patch_embed.num_patches,
|
|
embed_dim=encoder.embed_dim,
|
|
predictor_embed_dim=pred_emb_dim,
|
|
depth=pred_depth,
|
|
num_heads=encoder.num_heads)
|
|
|
|
def init_weights(m):
|
|
if isinstance(m, torch.nn.Linear):
|
|
trunc_normal_(m.weight, std=0.02)
|
|
if m.bias is not None:
|
|
torch.nn.init.constant_(m.bias, 0)
|
|
elif isinstance(m, torch.nn.LayerNorm):
|
|
torch.nn.init.constant_(m.bias, 0)
|
|
torch.nn.init.constant_(m.weight, 1.0)
|
|
|
|
for m in encoder.modules():
|
|
init_weights(m)
|
|
|
|
for m in predictor.modules():
|
|
init_weights(m)
|
|
|
|
encoder.to(device)
|
|
predictor.to(device)
|
|
logger.info(encoder)
|
|
return encoder, predictor
|
|
|
|
|
|
def init_opt(
|
|
encoder,
|
|
predictor,
|
|
iterations_per_epoch,
|
|
start_lr,
|
|
ref_lr,
|
|
warmup,
|
|
num_epochs,
|
|
wd=1e-6,
|
|
final_wd=1e-6,
|
|
final_lr=0.0,
|
|
use_bfloat16=False,
|
|
ipe_scale=1.25
|
|
):
|
|
param_groups = [
|
|
{
|
|
'params': (p for n, p in encoder.named_parameters()
|
|
if ('bias' not in n) and (len(p.shape) != 1))
|
|
}, {
|
|
'params': (p for n, p in predictor.named_parameters()
|
|
if ('bias' not in n) and (len(p.shape) != 1))
|
|
}, {
|
|
'params': (p for n, p in encoder.named_parameters()
|
|
if ('bias' in n) or (len(p.shape) == 1)),
|
|
'WD_exclude': True,
|
|
'weight_decay': 0
|
|
}, {
|
|
'params': (p for n, p in predictor.named_parameters()
|
|
if ('bias' in n) or (len(p.shape) == 1)),
|
|
'WD_exclude': True,
|
|
'weight_decay': 0
|
|
}
|
|
]
|
|
|
|
logger.info('Using AdamW')
|
|
optimizer = torch.optim.AdamW(param_groups)
|
|
scheduler = WarmupCosineSchedule(
|
|
optimizer,
|
|
warmup_steps=int(warmup*iterations_per_epoch),
|
|
start_lr=start_lr,
|
|
ref_lr=ref_lr,
|
|
final_lr=final_lr,
|
|
T_max=int(ipe_scale*num_epochs*iterations_per_epoch))
|
|
wd_scheduler = CosineWDSchedule(
|
|
optimizer,
|
|
ref_wd=wd,
|
|
final_wd=final_wd,
|
|
T_max=int(ipe_scale*num_epochs*iterations_per_epoch))
|
|
scaler = torch.cuda.amp.GradScaler() if use_bfloat16 else None
|
|
return optimizer, scaler, scheduler, wd_scheduler
|