import math

import torch

import logging

import numpy as np

from tqdm import tqdm

import torch.nn as nn

import multiprocessing

from os.path import join

from datetime import datetime

import torchvision.transforms as transforms

from torch.utils.data.dataloader import DataLoader

torch.backends.cudnn.benchmark= True # Provides a speedup

import util

import test

import parser

import commons

import datasets_ws

import network

import warnings

warnings.filterwarnings('ignore')

from loss import LocalFeatureLoss

import os

os.environ["CUDA_VISIBLE_DEVICES"] = "0"

args = parser.parse_arguments()

start_time = datetime.now()

args.save_dir = join("logs", args.save_dir, start_time.strftime('%Y-%m-%d_%H-%M-%S'))

commons.setup_logging(args.save_dir)

commons.make_deterministic(args.seed)

logging.info(f"Arguments: {args}")

logging.info(f"The outputs are being saved in {args.save_dir}")

logging.info(f"Using {torch.cuda.device_count()} GPUs and {multiprocessing.cpu_count()} CPUs")

logging.debug(f"Loading dataset {args.dataset_name} from folder {args.datasets_folder}")

triplets_ds = datasets_ws.TripletsDataset(args, args.datasets_folder, args.dataset_name, "train", args.negs_num_per_query)

logging.info(f"Train query set: {triplets_ds}")

val_ds = datasets_ws.BaseDataset(args, args.datasets_folder, args.dataset_name, "val")

logging.info(f"Val set: {val_ds}")

test_ds = datasets_ws.BaseDataset(args, args.datasets_folder, args.dataset_name, "test")

logging.info(f"Test set: {test_ds}")

model = network.GeoLocalizationNet(args)

model = model.to(args.device)

model = torch.nn.DataParallel(model)

for name, param in model.module.backbone.named_parameters():

if "adapter" not in name:

param.requires_grad = False

for n, m in model.named_modules():

if 'adapter' in n:

for n2, m2 in m.named_modules():

if 'D_fc2' in n2:

if isinstance(m2, nn.Linear):

nn.init.constant_(m2.weight, 0.)

nn.init.constant_(m2.bias, 0.)

if args.optim == "adam":

optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)

# scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=int(args.queries_per_epoch/args.train_batch_size), gamma=0.5, last_epoch=-1)

elif args.optim == "sgd":

optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=0.9, weight_decay=0.001)

GlobalTriplet = nn.TripletMarginLoss(margin=args.margin, p=2, reduction="sum")

MNNLocalFeatureLoss = LocalFeatureLoss().to(args.device)

if args.resume:

model, _, best_r5, start_epoch_num, not_improved_num = util.resume_train(args, model, strict=False)

logging.info(f"Resuming from epoch {start_epoch_num} with best recall@5 {best_r5:.1f}")

else:

best_r5 = start_epoch_num = not_improved_num = 0

logging.info(f"Output dimension of the model is {args.features_dim}")

for epoch_num in range(start_epoch_num, args.epochs_num):

logging.info(f"Start training epoch: {epoch_num:02d}")

epoch_start_time = datetime.now()

epoch_losses = np.zeros((0,1), dtype=np.float32)

# How many loops should an epoch last (default is 5000/1000=5)

loops_num = math.ceil(args.queries_per_epoch / args.cache_refresh_rate)

for loop_num in range(loops_num):

logging.debug(f"Cache: {loop_num} / {loops_num}")

# Compute triplets to use in the triplet loss

triplets_ds.is_inference = True

triplets_ds.compute_triplets(args, model)

triplets_ds.is_inference = False

triplets_dl = DataLoader(dataset=triplets_ds, num_workers=args.num_workers,

batch_size=args.train_batch_size,

collate_fn=datasets_ws.collate_fn,

pin_memory=(args.device=="cuda"),

drop_last=True)

model = model.train()

# images shape: (train_batch_size*4)*3*H*W

for images, triplets_local_indexes, _ in tqdm(triplets_dl, ncols=100):

# Flip all triplets or none

if args.horizontal_flip:

images = transforms.RandomHorizontalFlip()(images)

# Compute features of all images (images contains queries, positives and negatives)

local_features, global_features = model(images.to(args.device))

total_loss = 0

global_loss = 0

local_loss = 0

if args.criterion == "triplet":

triplets_local_indexes = torch.transpose(

triplets_local_indexes.view(args.train_batch_size, args.negs_num_per_query, 3), 1, 0)

for triplets in triplets_local_indexes:

queries_indexes, positives_indexes, negatives_indexes = triplets.T

global_loss += GlobalTriplet(global_features[queries_indexes],

global_features[positives_indexes],

global_features[negatives_indexes])

local_loss += MNNLocalFeatureLoss([local_features[queries_indexes],

local_features[positives_indexes],

local_features[negatives_indexes]])

global_loss /= (args.train_batch_size * args.negs_num_per_query)

local_loss /= (args.train_batch_size * args.negs_num_per_query)

total_loss = global_loss + local_loss

del global_features

del local_features

optimizer.zero_grad()

total_loss.backward()

optimizer.step()

batch_loss = total_loss.item()

epoch_losses = np.append(epoch_losses, batch_loss)

del total_loss

logging.debug(f"global loss = {global_loss.item():.6f}, " +

f"local loss = {local_loss.item():.6f}")

logging.debug(f"Epoch[{epoch_num:02d}]({loop_num}/{loops_num}): " +

f"current batch triplet loss = {batch_loss:.4f}, " +

f"average epoch triplet loss = {epoch_losses.mean():.4f}")

logging.info(f"Finished epoch {epoch_num:02d} in {str(datetime.now() - epoch_start_time)[:-7]}, "

f"average epoch triplet loss = {epoch_losses.mean():.4f}")

# Compute recalls on validation set

recalls, recalls_str = test.test(args, val_ds, model)

logging.info(f"Reranking recalls on val set {val_ds}: {recalls_str}")

is_best = recalls[1] > best_r5

# Save checkpoint, which contains all training parameters

util.save_checkpoint(args, {"epoch_num": epoch_num, "model_state_dict": model.state_dict(),

"optimizer_state_dict": optimizer.state_dict(), "recalls": recalls, "best_r5": best_r5,

"not_improved_num": not_improved_num

}, is_best, filename="last_model.pth")

# If recall@5 did not improve for "many" epochs, stop training

if is_best:

logging.info(f"Improved: previous best R@5 = {best_r5:.1f}, current R@5 = {(recalls[1]):.1f}")

best_r5 = recalls[1]

not_improved_num = 0

else:

not_improved_num += 1

logging.info(f"Not improved: {not_improved_num} / {args.patience}: best R@5 = {best_r5:.1f}, current R@5 = {(recalls[1]):.1f}")

if not_improved_num >= args.patience:

logging.info(f"Performance did not improve for {not_improved_num} epochs. Stop training.")

break

logging.info(f"Best R@5: {best_r5:.1f}")

logging.info(f"Trained for {epoch_num+1:02d} epochs, in total in {str(datetime.now() - start_time)[:-7]}")

best_model_state_dict = torch.load(join(args.save_dir, "best_model.pth"))["model_state_dict"]

model.load_state_dict(best_model_state_dict)

recalls, recalls_str = test.test(args, test_ds, model, test_method=args.test_method)

logging.info(f"Recalls on {test_ds}: {recalls_str}")