Pikaliov
4ff36ce188
Add dataloaders (v1/v2), analysis scripts, and documentation for working with UAV-GeoLoc (World-UAV). Co-authored-by: Cursor <cursoragent@cursor.com>
341 lines
14 KiB
Python
341 lines
14 KiB
Python
import os
|
|
import time
|
|
import numpy as np
|
|
import math
|
|
import shutil
|
|
import sys
|
|
import torch
|
|
from dataclasses import dataclass,field
|
|
from torch.cuda.amp import GradScaler
|
|
from torch.utils.data import DataLoader
|
|
from transformers import get_constant_schedule_with_warmup, get_polynomial_decay_schedule_with_warmup, get_cosine_schedule_with_warmup
|
|
from torchvision import transforms as T
|
|
from torch.utils.tensorboard import SummaryWriter
|
|
|
|
from dataset.World import WorldDatasetTrainVanilia, WorldDatasetEvalVanilia
|
|
from models import model,trainer
|
|
from utils import setting
|
|
from utils import loss
|
|
from eval import eval
|
|
|
|
|
|
|
|
def default_backbone_config():
|
|
|
|
return {
|
|
"backbone_arch" : "resnet18",
|
|
"pretrain_flag":True
|
|
}
|
|
|
|
def default_agg_config():
|
|
|
|
return {
|
|
"agg_arch": "multiconvap", #convap
|
|
"agg_config": {
|
|
"in_channels": 512, #256 #512
|
|
"out_channels": 512, #256
|
|
"s1": 1,
|
|
"s2": 1,
|
|
'LPN':False
|
|
}
|
|
}
|
|
|
|
@dataclass
|
|
class Configuration:
|
|
|
|
model: str = "resnet-new-all-frozen"
|
|
|
|
# Savepath for model checkpoints
|
|
model_path: str = "./world_vanilia"
|
|
|
|
# model config
|
|
backbone:dict = field(default_factory=default_backbone_config)
|
|
|
|
agg:dict = field(default_factory=default_agg_config)
|
|
|
|
# dataset
|
|
dataset_root_dir: str = "/media/Shen/Data/RingoData/WorldLoc"
|
|
train_query_txt: str = "/media/Shen/Data/RingoData/WorldLoc/Index/train_query_all.txt"
|
|
|
|
# val_index
|
|
val_index_txt = "/media/Shen/Data/RingoData/WorldLoc/Index/val_all.txt"
|
|
|
|
# test_index
|
|
test_index_txt = "/media/Shen/Data/RingoData/WorldLoc/Index/test_country.txt"
|
|
|
|
|
|
# Checkpoint to start from
|
|
checkpoint_start = None
|
|
|
|
# set num_workers to 0 if on Windows
|
|
num_workers: int = 0 if os.name == 'nt' else 4
|
|
|
|
# train on GPU if available
|
|
device: str = 'cuda' if torch.cuda.is_available() else 'cpu'
|
|
|
|
# for better performance
|
|
cudnn_benchmark: bool = True
|
|
|
|
# make cudnn deterministic
|
|
cudnn_deterministic: bool = False
|
|
|
|
# trainning
|
|
mixed_precision: bool = True
|
|
custom_sampling: bool = True # use custom sampling instead of random
|
|
seed = 1
|
|
epochs: int = 10
|
|
batch_size: int = 128 # keep in mind real_batch_size = 2 * batch_size 128
|
|
verbose: bool = True
|
|
gpu_ids: tuple = (0,2,3) # GPU ids for training
|
|
|
|
# Optimizer
|
|
clip_grad = 100. # None | float
|
|
decay_exclue_bias: bool = False
|
|
grad_checkpointing: bool = False # Gradient Checkpointing
|
|
|
|
# Loss
|
|
label_smoothing: float = 0.1
|
|
|
|
# Learning Rate
|
|
lr: float = 0.001 # 1 * 10^-4 for ViT | 1 * 10^-1 for CNN
|
|
scheduler: str = "cosine" # "polynomial" | "cosine" | "constant" | None
|
|
warmup_epochs: int = 0.1
|
|
lr_end: float = 0.0001 # only for "polynomial"
|
|
|
|
#-------------------------------------------------------------------------------------------#
|
|
# Train Config
|
|
#-------------------------------------------------------------------------------------------#
|
|
config = Configuration()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
model_path = "{}/{}/{}".format(config.model_path,
|
|
config.model,
|
|
time.strftime("%H%M%S"))
|
|
|
|
if not os.path.exists(model_path):
|
|
os.makedirs(model_path)
|
|
shutil.copyfile(os.path.basename(__file__), "{}/train.py".format(model_path))
|
|
|
|
# Redirect print to both console and log file
|
|
sys.stdout = setting.Logger(os.path.join(model_path, 'log.txt'))
|
|
|
|
setting.setup_system(seed=config.seed,
|
|
cudnn_benchmark=config.cudnn_benchmark,
|
|
cudnn_deterministic=config.cudnn_deterministic)
|
|
|
|
#-----------------------------------------------------------------------------#
|
|
# Model #
|
|
#-----------------------------------------------------------------------------#
|
|
|
|
print("\nModel: {}".format(config.model))
|
|
|
|
|
|
# backbone
|
|
model = model.BackboneGlobal(config.backbone['backbone_arch'],
|
|
config.backbone['pretrain_flag'],
|
|
config.agg['agg_arch'],
|
|
config.agg['agg_config'])
|
|
|
|
# Load pretrained Checkpoint
|
|
if config.checkpoint_start is not None:
|
|
print("Start from:", config.checkpoint_start)
|
|
model_state_dict = torch.load(config.checkpoint_start)
|
|
model.load_state_dict(model_state_dict, strict=False)
|
|
|
|
# Data parallel
|
|
print("GPUs available:", torch.cuda.device_count())
|
|
if torch.cuda.device_count() > 1 and len(config.gpu_ids) > 1:
|
|
model = torch.nn.DataParallel(model, device_ids=config.gpu_ids)
|
|
|
|
# Model to device
|
|
model = model.to(config.device)
|
|
|
|
#------------------------setting dataset-------------------------------------------------#
|
|
IMAGENET_MEAN_STD = {'mean': [0.485, 0.456, 0.406],
|
|
'std': [0.229, 0.224, 0.225]}
|
|
train_transform = T.Compose([
|
|
T.Resize((224, 224), interpolation=T.InterpolationMode.BILINEAR),
|
|
T.RandAugment(num_ops=3, interpolation=T.InterpolationMode.BILINEAR),
|
|
T.AugMix(),
|
|
# T.ColorJitter(brightness=0.5, contrast=0.1, saturation=0.1,
|
|
# hue=0),
|
|
# T.RandomGrayscale(p=0.2),
|
|
# T.RandomPosterize(p=0.2, bits=4),
|
|
# T.GaussianBlur(kernel_size=(1, 5), sigma=(0.1, 5)),
|
|
T.ToTensor(),
|
|
T.Normalize(mean=IMAGENET_MEAN_STD["mean"], std=IMAGENET_MEAN_STD["std"]),
|
|
])
|
|
|
|
eval_transform = T.Compose([
|
|
T.Resize((224, 224), interpolation=T.InterpolationMode.BILINEAR),
|
|
T.ToTensor(),
|
|
T.Normalize(mean=IMAGENET_MEAN_STD["mean"], std=IMAGENET_MEAN_STD["std"]),
|
|
])
|
|
|
|
#-----------------------------------------------------------------------------#
|
|
# DataLoader #
|
|
#-----------------------------------------------------------------------------#
|
|
|
|
train_dataset = WorldDatasetTrainVanilia(data_dir=config.dataset_root_dir,
|
|
query_txt=config.train_query_txt,
|
|
transforms_query=train_transform,
|
|
transforms_db=train_transform,
|
|
shuffle_batch_size=config.batch_size)
|
|
|
|
|
|
train_dataloader = DataLoader(train_dataset,
|
|
batch_size=config.batch_size,
|
|
num_workers=config.num_workers,
|
|
shuffle=config.custom_sampling,
|
|
pin_memory=True)
|
|
|
|
#-----------------------------------------------------------------------------#
|
|
# Loss #
|
|
#-----------------------------------------------------------------------------#
|
|
|
|
# InfoNCE loss
|
|
loss_fn = torch.nn.CrossEntropyLoss(label_smoothing=config.label_smoothing)
|
|
loss_function = loss.InfoNCE(loss_function=loss_fn,
|
|
device=config.device,
|
|
)
|
|
# Supervised Contrastive loss
|
|
# loss_function = loss.SupervisedContrastiveLoss(temperature = 0.07, device=config.device)
|
|
|
|
if config.mixed_precision:
|
|
scaler = GradScaler(init_scale=2.**10)
|
|
else:
|
|
scaler = None
|
|
|
|
#-----------------------------------------------------------------------------#
|
|
# optimizer #
|
|
#-----------------------------------------------------------------------------#
|
|
|
|
optimizer = torch.optim.AdamW(model.parameters(), lr=config.lr)
|
|
|
|
#-----------------------------------------------------------------------------#
|
|
# Scheduler #
|
|
#-----------------------------------------------------------------------------#
|
|
|
|
train_steps = len(train_dataloader) * config.epochs
|
|
warmup_steps = len(train_dataloader) * config.warmup_epochs
|
|
|
|
if config.scheduler == "polynomial":
|
|
print("\nScheduler: polynomial - max LR: {} - end LR: {}".format(config.lr, config.lr_end))
|
|
scheduler = get_polynomial_decay_schedule_with_warmup(optimizer,
|
|
num_training_steps=train_steps,
|
|
lr_end = config.lr_end,
|
|
power=1.5,
|
|
num_warmup_steps=warmup_steps)
|
|
|
|
elif config.scheduler == "cosine":
|
|
print("\nScheduler: cosine - max LR: {}".format(config.lr))
|
|
scheduler = get_cosine_schedule_with_warmup(optimizer,
|
|
num_training_steps=train_steps,
|
|
num_warmup_steps=warmup_steps)
|
|
|
|
elif config.scheduler == "constant":
|
|
print("\nScheduler: constant - max LR: {}".format(config.lr))
|
|
scheduler = get_constant_schedule_with_warmup(optimizer,
|
|
num_warmup_steps=warmup_steps)
|
|
|
|
else:
|
|
scheduler = None
|
|
|
|
print("Warmup Epochs: {} - Warmup Steps: {}".format(str(config.warmup_epochs).ljust(2), warmup_steps))
|
|
print("Train Epochs: {} - Train Steps: {}".format(config.epochs, train_steps))
|
|
|
|
|
|
#-----------------------------------------------------------------------------#
|
|
# Shuffle #
|
|
#-----------------------------------------------------------------------------#
|
|
if config.custom_sampling:
|
|
train_dataloader.dataset.shuffle()
|
|
|
|
#-----------------------------------------------------------------------------#
|
|
# Train #
|
|
#-----------------------------------------------------------------------------#
|
|
start_epoch = 0
|
|
best_score = 0
|
|
|
|
#-----------------------------------------------------------------------------#
|
|
# Writer
|
|
#-----------------------------------------------------------------------------#
|
|
# Writer
|
|
writer = SummaryWriter('world_vanillia/cnn' + config.model)
|
|
LPN_flag = config.agg['agg_config']['LPN']
|
|
|
|
|
|
for epoch in range(1, config.epochs+1):
|
|
|
|
print("\n{}[Epoch: {}]{}".format(30*"-", epoch, 30*"-"))
|
|
|
|
|
|
train_loss = trainer.train_backbone(config,
|
|
model,
|
|
dataloader=train_dataloader,
|
|
loss_function=loss_function,
|
|
optimizer=optimizer,
|
|
scheduler=scheduler,
|
|
scaler=scaler,
|
|
writer=writer,
|
|
LPN=LPN_flag)
|
|
|
|
print("Epoch: {}, Train Loss = {:.3f}, Lr = {:.6f}".format(epoch,
|
|
train_loss,
|
|
optimizer.param_groups[0]['lr']))
|
|
|
|
#------------------------------------------------------------Eval---------------------------------------------------------------------#
|
|
result_list = []
|
|
with open(config.val_index_txt,"r") as val_test:
|
|
for line in val_test:
|
|
eva_dataset_query = WorldDatasetEvalVanilia(data_dir=config.dataset_root_dir,
|
|
name=line.strip('\n'),
|
|
mode='query',
|
|
transforms=eval_transform)
|
|
|
|
eval_dataloader_query = DataLoader(eva_dataset_query,
|
|
batch_size=config.batch_size,
|
|
num_workers=config.num_workers,
|
|
shuffle=not config.custom_sampling,
|
|
pin_memory=True)
|
|
|
|
eva_dataset_db = WorldDatasetEvalVanilia(data_dir=config.dataset_root_dir,
|
|
name=line.strip('\n'),
|
|
mode='DB',
|
|
transforms=eval_transform)
|
|
|
|
eval_dataloader_db = DataLoader(eva_dataset_db,
|
|
batch_size=config.batch_size,
|
|
num_workers=config.num_workers,
|
|
shuffle=not config.custom_sampling,
|
|
pin_memory=True)
|
|
|
|
pos_gt = eval_dataloader_db.dataset.get_gt()
|
|
result,_ , _ = eval.evaluate(config, model, eval_dataloader_query, eval_dataloader_db, pos_gt, mode='vanilia',LPN=config.agg['agg_config']['LPN'])
|
|
print(line.strip('\n'), 'top 1: ', round(result[0]*100,2), 'top 5: ', round(result[1]*100,2), 'top 10: ', round(result[2]*100,2))
|
|
result_list.append(result)
|
|
writer.add_scalar(line.strip('\n'), round(result[0]*100,2), epoch)
|
|
|
|
|
|
result_array = np.array(result_list)
|
|
average_result = np.mean(result_array, axis=0)
|
|
print('Average', 'top 1: ', round(average_result[0]*100,2), 'top 5: ', round(average_result[1]*100,2), 'top 10: ', round(average_result[2]*100,2))
|
|
writer.add_scalar('Average/top1', round(average_result[0]*100,2), epoch)
|
|
writer.add_scalar('Average/top5', round(average_result[1]*100,2), epoch)
|
|
|
|
#------------------------------------------------------------Save---------------------------------------------------------------------#
|
|
if average_result[0] > best_score:
|
|
|
|
best_score = average_result[0]
|
|
|
|
if torch.cuda.device_count() > 1 and len(config.gpu_ids) > 1:
|
|
torch.save(model.module.state_dict(), '{}/weights_e{}_{:.4f}.pth'.format(model_path, epoch, average_result[0]))
|
|
else:
|
|
torch.save(model.state_dict(), '{}/weights_e{}_{:.4f}.pth'.format(model_path, epoch, average_result[0]))
|
|
|
|
|
|
if config.custom_sampling:
|
|
train_dataloader.dataset.shuffle()
|