import argparse import datetime import json import numpy as np import os import time import random from pathlib import Path import sys from PIL import Image import torch.nn.functional as F import torch import torch.backends.cudnn as cudnn from torch.utils.data import Dataset import torchvision import wandb import timm from tqdm import tqdm assert "0.4.5" <= timm.__version__ <= "0.4.9" # version check import timm.optim.optim_factory as optim_factory import util.misc as misc from util.misc import NativeScalerWithGradNormCount as NativeScaler import util.lr_sched as lr_sched from util.FSC147 import transform_train, transform_val import models_mae_cross def get_args_parser(): parser = argparse.ArgumentParser('MAE pre-training', add_help=True) parser.add_argument('--batch_size', default=26, type=int, help='Batch size per GPU (effective batch size is batch_size * accum_iter * # gpus)') parser.add_argument('--epochs', default=200, type=int) parser.add_argument('--accum_iter', default=1, type=int, help='Accumulate gradient iterations (for increasing the effective batch size under memory constraints)') # Model parameters parser.add_argument('--model', default='mae_vit_base_patch16', type=str, metavar='MODEL', help='Name of model to train') parser.add_argument('--mask_ratio', default=0.5, type=float, help='Masking ratio (percentage of removed patches).') parser.add_argument('--norm_pix_loss', action='store_true', help='Use (per-patch) normalized pixels as targets for computing loss') parser.set_defaults(norm_pix_loss=False) # Optimizer parameters parser.add_argument('--weight_decay', type=float, default=0.05, help='weight decay (default: 0.05)') parser.add_argument('--lr', type=float, default=None, metavar='LR', help='learning rate (absolute lr)') parser.add_argument('--blr', type=float, default=1e-3, metavar='LR', help='base learning rate: absolute_lr = base_lr * total_batch_size / 256') parser.add_argument('--min_lr', type=float, default=0., metavar='LR', help='lower lr bound for cyclic schedulers that hit 0') parser.add_argument('--warmup_epochs', type=int, default=10, metavar='N', help='epochs to warmup LR') # Dataset parameters parser.add_argument('--data_path', default='./data/FSC147/', type=str, help='dataset path') parser.add_argument('--anno_file', default='annotation_FSC147_pos.json', type=str, help='annotation json file for positive samples') parser.add_argument('--anno_file_negative', default='./data/FSC147/annotation_FSC147_neg.json', type=str, help='annotation json file for negative samples') parser.add_argument('--data_split_file', default='Train_Test_Val_FSC_147.json', type=str, help='data split json file') parser.add_argument('--class_file', default='ImageClasses_FSC147.txt', type=str, help='class json file') parser.add_argument('--im_dir', default='images_384_VarV2', type=str, help='images directory') parser.add_argument('--output_dir', default='./data/out/fim6_dir', help='path where to save, empty for no saving') parser.add_argument('--device', default='cuda', help='device to use for training / testing') parser.add_argument('--seed', default=0, type=int) parser.add_argument('--resume', default='./data/checkpoint.pth', help='resume from checkpoint') parser.add_argument('--do_resume', action='store_true', help='Resume training (e.g. if crashed).') # Training parameters parser.add_argument('--start_epoch', default=0, type=int, metavar='N', help='start epoch') parser.add_argument('--num_workers', default=10, type=int) parser.add_argument('--pin_mem', action='store_true', help='Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.') parser.add_argument('--no_pin_mem', action='store_false', dest='pin_mem') parser.set_defaults(pin_mem=True) parser.add_argument('--do_aug', action='store_true', help='Perform data augmentation.') parser.add_argument('--no_do_aug', action='store_false', dest='do_aug') parser.set_defaults(do_aug=True) # Distributed training parameters parser.add_argument('--world_size', default=1, type=int, help='number of distributed processes') parser.add_argument('--local_rank', default=-1, type=int) parser.add_argument('--dist_on_itp', action='store_true') parser.add_argument('--dist_url', default='env://', help='url used to set up distributed training') # Logging parameters parser.add_argument("--title", default="count", type=str) parser.add_argument("--wandb", default="240227", type=str) parser.add_argument("--team", default="wsense", type=str) parser.add_argument("--wandb_id", default=None, type=str) return parser os.environ["CUDA_LAUNCH_BLOCKING"] = '0' class TrainData(Dataset): def __init__(self, args, split='train', do_aug=True): with open(args.anno_file) as f: annotations = json.load(f) # Load negative annotations with open(args.anno_file_negative) as f: neg_annotations = json.load(f) with open(args.data_split_file) as f: data_split = json.load(f) self.img = data_split[split] random.shuffle(self.img) self.split = split self.img_dir = im_dir self.TransformTrain = transform_train(args, do_aug=do_aug) self.TransformVal = transform_val(args) self.annotations = annotations self.neg_annotations = neg_annotations self.im_dir = im_dir def __len__(self): return len(self.img) def __getitem__(self, idx): im_id = self.img[idx] anno = self.annotations[im_id] bboxes = anno['box_examples_coordinates'] dots = np.array(anno['points']) # 加载负样本的框 neg_anno = self.neg_annotations[im_id] # 假设每个图像ID在负样本注释中都有对应的条目 neg_bboxes = neg_anno['box_examples_coordinates'] rects = list() for bbox in bboxes: x1 = bbox[0][0] y1 = bbox[0][1] x2 = bbox[2][0] y2 = bbox[2][1] if x1 < 0: x1 = 0 if x2 < 0: x2 = 0 if y1 < 0: y1 = 0 if y2 < 0: y2 = 0 rects.append([y1, x1, y2, x2]) neg_rects = list() for neg_bbox in neg_bboxes: x1 = neg_bbox[0][0] y1 = neg_bbox[0][1] x2 = neg_bbox[2][0] y2 = neg_bbox[2][1] if x1 < 0: x1 = 0 if x2 < 0: x2 = 0 if y1 < 0: y1 = 0 if y2 < 0: y2 = 0 neg_rects.append([y1, x1, y2, x2]) image = Image.open('{}/{}'.format(self.im_dir, im_id)) if image.mode == "RGBA": image = image.convert("RGB") image.load() m_flag = 0 sample = {'image': image, 'lines_boxes': rects, 'neg_lines_boxes': neg_rects,'dots': dots, 'id': im_id, 'm_flag': m_flag} sample = self.TransformTrain(sample) if self.split == "train" else self.TransformVal(sample) return sample['image'], sample['gt_density'], len(dots), sample['boxes'],sample['neg_boxes'], sample['pos'],sample['m_flag'], im_id def main(args): wandb_run = None try: misc.init_distributed_mode(args) print('job dir: {}'.format(os.path.dirname(os.path.realpath(__file__)))) print("{}".format(args).replace(', ', ',\n')) device = torch.device(args.device) # if torch.cuda.is_available(): # device = torch.device("cuda:5") # fix the seed for reproducibility seed = args.seed + misc.get_rank() torch.manual_seed(seed) np.random.seed(seed) cudnn.benchmark = True dataset_train = TrainData(args, do_aug=args.do_aug) dataset_val = TrainData(args, split='val') num_tasks = misc.get_world_size() global_rank = misc.get_rank() sampler_train = torch.utils.data.DistributedSampler( dataset_train, num_replicas=num_tasks, rank=global_rank, shuffle=True ) sampler_val = torch.utils.data.DistributedSampler( dataset_val, num_replicas=num_tasks, rank=global_rank, shuffle=True ) if global_rank == 0: if args.wandb is not None: wandb_run = wandb.init( config=args, resume="allow", project=args.wandb, name=args.title, # entity=args.team, tags=["count", "finetuning"], id=args.wandb_id, ) data_loader_train = torch.utils.data.DataLoader( dataset_train, sampler=sampler_train, batch_size=args.batch_size, num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=False, ) data_loader_val = torch.utils.data.DataLoader( dataset_val, sampler=sampler_val, batch_size=args.batch_size, num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=False, ) # define the model model = models_mae_cross.__dict__[args.model](norm_pix_loss=args.norm_pix_loss) model.to(device) model_without_ddp = model # print("Model = %s" % str(model_without_ddp)) eff_batch_size = args.batch_size * args.accum_iter * misc.get_world_size() if args.lr is None: # only base_lr is specified args.lr = args.blr * eff_batch_size / 256 print("base lr: %.2e" % (args.lr * 256 / eff_batch_size)) print("actual lr: %.2e" % args.lr) print("accumulate grad iterations: %d" % args.accum_iter) print("effective batch size: %d" % eff_batch_size) if args.distributed: model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu], find_unused_parameters=True) model_without_ddp = model.module # following timm: set wd as 0 for bias and norm layers param_groups = optim_factory.add_weight_decay(model_without_ddp, args.weight_decay) optimizer = torch.optim.AdamW(param_groups, lr=args.lr, betas=(0.9, 0.95)) print(optimizer) loss_scaler = NativeScaler() min_MAE = 99999 print_freq = 50 save_freq = 50 misc.load_model_FSC_full(args=args, model_without_ddp=model_without_ddp, optimizer=optimizer, loss_scaler=loss_scaler) print(f"Start training for {args.epochs - args.start_epoch} epochs - rank {global_rank}") start_time = time.time() for epoch in range(args.start_epoch, args.epochs): if args.distributed: data_loader_train.sampler.set_epoch(epoch) # train one epoch model.train(True) accum_iter = args.accum_iter # some parameters in training train_mae = torch.tensor([0], dtype=torch.float64, device=device) train_mse = torch.tensor([0], dtype=torch.float64, device=device) val_mae = torch.tensor([0], dtype=torch.float64, device=device) val_mse = torch.tensor([0], dtype=torch.float64, device=device) val_nae = torch.tensor([0], dtype=torch.float64, device=device) optimizer.zero_grad() for data_iter_step, (samples, gt_density, _, pos_boxes, neg_boxes, pos, m_flag, im_names) in enumerate( tqdm(data_loader_train, total=len(data_loader_train), desc=f"Train [e. {epoch} - r. {global_rank}]")): idx = data_iter_step + (epoch * len(data_loader_train)) if data_iter_step % accum_iter == 0: lr_sched.adjust_learning_rate(optimizer, data_iter_step / len(data_loader_train) + epoch, args) samples = samples.to(device, non_blocking=True, dtype=torch.half) gt_density = gt_density.to(device, non_blocking=True, dtype=torch.half) pos_boxes = pos_boxes.to(device, non_blocking=True, dtype=torch.half) neg_boxes = neg_boxes.to(device, non_blocking=True, dtype=torch.half) # 如果至少有一个图像在批处理中使用了Type 2 Mosaic，则禁止0-shot。 flag = 0 for i in range(m_flag.shape[0]): flag += m_flag[i].item() if flag == 0: shot_num = random.randint(0, 3) else: shot_num = random.randint(1, 3) with torch.cuda.amp.autocast(): pos_output = model(samples, pos_boxes, shot_num) # 正样本输出 # 计算正样本损失 mask = np.random.binomial(n=1, p=0.8, size=[384, 384]) masks = np.tile(mask, (pos_output.shape[0], 1)) masks = masks.reshape(pos_output.shape[0], 384, 384) masks = torch.from_numpy(masks).to(device) pos_loss = ((pos_output - gt_density) ** 2) pos_loss = (pos_loss * masks / (384 * 384)).sum() / pos_output.shape[0] # 负样本输出 with torch.cuda.amp.autocast(): neg_output = model(samples, neg_boxes, 1) # 负样本输出 cnt1 = 1-torch.exp(-(torch.abs(pos_output.sum()/60 - gt_density.sum()/60).mean())) if neg_output.shape[0] == 0: cnt2 = 0 else: # cnt2 = torch.log(torch.abs((neg_output.sum() / neg_output.shape[0]) - 1).mean()+1) cnt2 = 1-torch.exp(-(torch.abs((neg_output.sum() / (neg_output.shape[0]*60)) - 1).mean())) cnt = cnt1+cnt2 # 计算正样本损失 mask = np.random.binomial(n=1, p=0.8, size=[384, 384]) masks = np.tile(mask, (neg_output.shape[0], 1)) masks = masks.reshape(neg_output.shape[0], 384, 384) masks = torch.from_numpy(masks).to(device) neg_loss = ((neg_output - gt_density) ** 2) if neg_output.shape[0] == 0: neg_loss = 1 else: neg_loss = (neg_loss * masks / (384 * 384)).sum() / neg_output.shape[0] margin = 0.5 contrastive_loss = torch.relu(pos_loss - neg_loss + margin) total_loss = contrastive_loss+pos_loss # 更新 MAE 和 RMSE with torch.no_grad(): pred_cnt = (pos_output.view(len(samples), -1)).sum(1) / 60 gt_cnt = (gt_density.view(len(samples), -1)).sum(1) / 60 cnt_err = torch.abs(pred_cnt - gt_cnt).float() batch_mae = cnt_err.double().mean() batch_mse = (cnt_err ** 2).double().mean() train_mae += batch_mae train_mse += batch_mse if not torch.isfinite(total_loss): print("Loss is {}, stopping training".format(total_loss)) sys.exit(1) total_loss /= accum_iter loss_scaler(total_loss, optimizer, parameters=model.parameters(), update_grad=(data_iter_step + 1) % accum_iter == 0) if (data_iter_step + 1) % accum_iter == 0: optimizer.zero_grad() lr = optimizer.param_groups[0]["lr"] loss_value_reduce = misc.all_reduce_mean(total_loss) if (data_iter_step + 1) % (print_freq * accum_iter) == 0 and (data_iter_step + 1) != len(data_loader_train) and data_iter_step != 0: if wandb_run is not None: log = {"train/loss": loss_value_reduce, "train/lr": lr, "train/MAE": batch_mae, "train/RMSE": batch_mse ** 0.5} wandb.log(log, step=idx) # evaluation on Validation split for val_samples, val_gt_density, val_n_ppl, val_boxes,_, val_pos, _, val_im_names in \ tqdm(data_loader_val, total=len(data_loader_val), desc=f"Val [e. {epoch} - r. {global_rank}]"): val_samples = val_samples.to(device, non_blocking=True, dtype=torch.half) val_gt_density = val_gt_density.to(device, non_blocking=True, dtype=torch.half) val_boxes = val_boxes.to(device, non_blocking=True, dtype=torch.half) val_n_ppl = val_n_ppl.to(device, non_blocking=True) shot_num = random.randint(0, 3) with torch.no_grad(): with torch.cuda.amp.autocast(): val_output = model(val_samples, val_boxes, shot_num) val_pred_cnt = (val_output.view(len(val_samples), -1)).sum(1) / 60 val_gt_cnt = (val_gt_density.view(len(val_samples), -1)).sum(1) / 60 # print('val_pred_cnt',val_pred_cnt) # print('val_gt_cnt',val_gt_cnt) val_cnt_err = torch.abs(val_pred_cnt - val_gt_cnt).float() # print('val_cnt_err',val_cnt_err.mean()) val_cnt_err[val_cnt_err == float('inf')] = 0 val_mae += val_cnt_err.double().mean() # val_mae += val_cnt_err # print('val_mae',val_mae.mean()) val_cnt_err[val_cnt_err == float('inf')] = 0 val_mse += (val_cnt_err ** 2).double().mean() # val_mse += (val_cnt_err ** 2) _val_nae = val_cnt_err / val_gt_cnt _val_nae[_val_nae == float('inf')] = 0 val_nae += _val_nae.double().mean() # val_mae = val_mae/len(data_loader_val) # val_mse = val_mse/len(data_loader_val) # print('val_mae',val_mae) # print('val_mse',val_mse) # Output visualisation information to W&B if wandb_run is not None: train_wandb_densities = [] train_wandb_bboxes = [] val_wandb_densities = [] val_wandb_bboxes = [] black = torch.zeros([384, 384], device=device) for i in range(pos_output.shape[0]): # gt and predicted density w_d_map = torch.stack([pos_output[i], black, black]) gt_map = torch.stack([gt_density[i], black, black]) box_map = misc.get_box_map(samples[i], pos[i], device) w_gt_density = samples[i] / 2 + gt_map + box_map w_d_map_overlay = samples[i] / 2 + w_d_map w_densities = torch.cat([w_gt_density, w_d_map, w_d_map_overlay], dim=2) w_densities = torch.clamp(w_densities, 0, 1) train_wandb_densities += [wandb.Image(torchvision.transforms.ToPILImage()(w_densities), caption=f"[E#{epoch}] {im_names[i]} ({torch.sum(gt_density[i]).item()}, {torch.sum(pos_output[i]).item()})")] # exemplars w_boxes = torch.cat([pos_boxes[i][x, :, :, :] for x in range(pos_boxes[i].shape[0])], 2) train_wandb_bboxes += [wandb.Image(torchvision.transforms.ToPILImage()(w_boxes), caption=f"[E#{epoch}] {im_names[i]}")] for i in range(val_output.shape[0]): # gt and predicted density w_d_map = torch.stack([val_output[i], black, black]) gt_map = torch.stack([val_gt_density[i], black, black]) box_map = misc.get_box_map(val_samples[i], val_pos[i], device) w_gt_density = val_samples[i] / 2 + gt_map + box_map w_d_map_overlay = val_samples[i] / 2 + w_d_map w_densities = torch.cat([w_gt_density, w_d_map, w_d_map_overlay], dim=2) w_densities = torch.clamp(w_densities, 0, 1) val_wandb_densities += [wandb.Image(torchvision.transforms.ToPILImage()(w_densities), caption=f"[E#{epoch}] {val_im_names[i]} ({torch.sum(val_gt_density[i]).item()}, {torch.sum(val_output[i]).item()})")] # exemplars w_boxes = torch.cat([val_boxes[i][x, :, :, :] for x in range(val_boxes[i].shape[0])], 2) val_wandb_bboxes += [wandb.Image(torchvision.transforms.ToPILImage()(w_boxes), caption=f"[E#{epoch}] {val_im_names[i]}")] log = {"train/loss": loss_value_reduce, "train/lr": lr, "train/MAE": batch_mae, "train/RMSE": batch_mse ** 0.5, "val/MAE": val_mae / len(data_loader_val), "val/RMSE": (val_mse / len(data_loader_val)) ** 0.5, "val/NAE": val_nae / len(data_loader_val), "train_densitss": train_wandb_densities, "val_densites": val_wandb_densities, "train_boxes": train_wandb_bboxes, "val_boxes": val_wandb_bboxes} wandb.log(log, step=idx) # save train status and model if args.output_dir and (epoch % save_freq == 0 or epoch + 1 == args.epochs) and epoch != 0: misc.save_model( args=args, model=model, model_without_ddp=model_without_ddp, optimizer=optimizer, loss_scaler=loss_scaler, epoch=epoch, suffix=f"finetuning_{epoch}", upload=epoch % 100 == 0) elif True: misc.save_model( args=args, model=model, model_without_ddp=model_without_ddp, optimizer=optimizer, loss_scaler=loss_scaler, epoch=epoch, suffix=f"finetuning_last", upload=False) if args.output_dir and val_mae / len(data_loader_val) < min_MAE: min_MAE = val_mae / len(data_loader_val) misc.save_model( args=args, model=model, model_without_ddp=model_without_ddp, optimizer=optimizer, loss_scaler=loss_scaler, epoch=epoch, suffix="finetuning_minMAE") print(f'[Train Epoch #{epoch}] - MAE: {train_mae.item() / len(data_loader_train):5.2f}, RMSE: {(train_mse.item() / len(data_loader_train)) ** 0.5:5.2f}', flush=True) print(f'[Val Epoch #{epoch}] - MAE: {val_mae.item() / len(data_loader_val):5.2f}, RMSE: {(val_mse.item() / len(data_loader_val)) ** 0.5:5.2f}, NAE: {val_nae.item() / len(data_loader_val):5.2f}', flush=True) total_time = time.time() - start_time total_time_str = str(datetime.timedelta(seconds=int(total_time))) print('Training time {}'.format(total_time_str)) finally: if wandb_run is not None: wandb.run.finish() if __name__ == '__main__': args = get_args_parser() args = args.parse_args() data_path = Path(args.data_path) anno_file = data_path / args.anno_file data_split_file = data_path / args.data_split_file im_dir = data_path / args.im_dir if args.do_aug: class_file = data_path / args.class_file else: class_file = None args.anno_file = anno_file args.data_split_file = data_split_file args.im_dir = im_dir args.class_file = class_file if args.output_dir: Path(args.output_dir).mkdir(parents=True, exist_ok=True) main(args)