import random import argparse import os import time import numpy as np import matplotlib.pyplot as plt from tqdm import tqdm import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import DataLoader from accelerate import Accelerator from models.transformer import Dasheng_Encoder from models.sed_decoder import Decoder, TSED_Wrapper from dataset.tsed import TSED_AS from dataset.tsed_val import TSED_Val from utils import load_yaml_with_includes, get_lr_scheduler, ConcatDatasetBatchSampler from utils.data_aug import frame_shift, mixup, time_mask, feature_transformation from val import val_psds def parse_args(): parser = argparse.ArgumentParser() # Config settings parser.add_argument('--config-name', type=str, default='configs/model.yml') # Training settings parser.add_argument("--amp", type=str, default='fp16') parser.add_argument('--epochs', type=int, default=20) parser.add_argument('--num-workers', type=int, default=8) parser.add_argument('--num-threads', type=int, default=1) parser.add_argument('--eval-every-step', type=int, default=5000) parser.add_argument('--save-every-step', type=int, default=5000) # parser.add_argument('--dataloader', type=str, default='EACaps') parser.add_argument("--logit-normal-indices", type=bool, default=False) # Log and random seed parser.add_argument('--random-seed', type=int, default=2024) parser.add_argument('--log-step', type=int, default=100) parser.add_argument('--log-dir', type=str, default='../logs/') parser.add_argument('--save-dir', type=str, default='../ckpts/') return parser.parse_args() def setup_directories(args, params): args.log_dir = os.path.join(args.log_dir, params['model_name']) + '/' args.save_dir = os.path.join(args.save_dir, params['model_name']) + '/' os.makedirs(args.log_dir, exist_ok=True) os.makedirs(args.save_dir, exist_ok=True) def set_device(args): torch.set_num_threads(args.num_threads) if torch.cuda.is_available(): args.device = 'cuda' torch.cuda.manual_seed_all(args.random_seed) torch.backends.cuda.matmul.allow_tf32 = True if torch.backends.cudnn.is_available(): torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False else: args.device = 'cpu' if __name__ == '__main__': args = parse_args() params = load_yaml_with_includes(args.config_name) set_device(args) setup_directories(args, params) random.seed(args.random_seed) torch.manual_seed(args.random_seed) # use accelerator for multi-gpu training accelerator = Accelerator(mixed_precision=args.amp, gradient_accumulation_steps=params['opt']['accumulation_steps'], step_scheduler_with_optimizer=False) train_set = TSED_AS(**params['data']['train_data']) train_loader = DataLoader(train_set, shuffle=True, batch_size=params['opt']['batch_size'], num_workers=args.num_workers) val_set = TSED_Val(**params['data']['val_data']) val_loader = DataLoader(val_set, num_workers=0, batch_size=1, shuffle=False) # test_set = TSED_Val(**params['data']['test_data']) # test_loader = DataLoader(val_set, num_workers=0, batch_size=1, shuffle=False) encoder = Dasheng_Encoder(**params['encoder']).to(accelerator.device) pretrained_url = 'https://zenodo.org/records/11511780/files/dasheng_base.pt?download=1' dump = torch.hub.load_state_dict_from_url(pretrained_url, map_location='cpu') model_parmeters = dump['model'] # pretrained_url = 'https://zenodo.org/records/13315686/files/dasheng_audioset_mAP497.pt?download=1' # dump = torch.hub.load_state_dict_from_url(pretrained_url, map_location='cpu') # model_parmeters = dump encoder.load_state_dict(model_parmeters) decoder = Decoder(**params['decoder']).to(accelerator.device) model = TSED_Wrapper(encoder, decoder, params['ft_blocks'], params['frozen_encoder']) print(f"Trainable Parameters: {sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.2f}M") # model.load_state_dict(torch.load('../ckpts/TSED_AS_filter/20000.0.pt', map_location='cpu')['model']) if params['frozen_encoder']: optimizer = torch.optim.AdamW( model.parameters(), lr=params['opt']['learning_rate'], weight_decay=params['opt']['weight_decay'], betas=(params['opt']['beta1'], params['opt']['beta2']), eps=params['opt']['adam_epsilon']) else: optimizer = torch.optim.AdamW( [ {'params': model.encoder.parameters(), 'lr': 0.1 * params['opt']['learning_rate']}, {'params': model.decoder.parameters(), 'lr': params['opt']['learning_rate']} ], weight_decay=params['opt']['weight_decay'], betas=(params['opt']['beta1'], params['opt']['beta2']), eps=params['opt']['adam_epsilon']) lr_scheduler = get_lr_scheduler(optimizer, 'customized', **params['opt']['lr_scheduler']) strong_loss_func = nn.BCEWithLogitsLoss() model, optimizer, lr_scheduler, train_loader, val_loader = accelerator.prepare( model, optimizer, lr_scheduler, train_loader, val_loader) global_step = 0.0 losses = 0.0 if accelerator.is_main_process: model_module = model.module if hasattr(model, 'module') else model val_psds(model_module, val_loader, params, epoch='debug', split='val', save_path=args.log_dir + 'output/', device=accelerator.device) for epoch in range(args.epochs): model.train() for step, batch in enumerate(tqdm(train_loader)): with accelerator.accumulate(model): audio, cls, label, _ = batch mel = model.forward_to_spec(audio) # data aug mel, label = frame_shift(mel, label, params['net_pooling']) mel, label = time_mask(mel, label, params["net_pooling"], mask_ratios=params['data_aug']["time_mask_ratios"]) mel, _ = feature_transformation(mel, **params['data_aug']["transform"]) strong_pred = model(mel, cls) B, N, L = label.shape label = label.reshape(B * N, L) label = label.unsqueeze(1) loss = strong_loss_func(strong_pred, label) accelerator.backward(loss) # clip grad up if accelerator.sync_gradients: if 'grad_clip' in params['opt'] and params['opt']['grad_clip'] > 0: accelerator.clip_grad_norm_(model.parameters(), max_norm=params['opt']['grad_clip']) optimizer.step() lr_scheduler.step() optimizer.zero_grad() global_step += 1/params['opt']['accumulation_steps'] losses += loss.item()/params['opt']['accumulation_steps'] if accelerator.is_main_process: if global_step % args.log_step == 0: current_time = time.asctime(time.localtime(time.time())) epoch_info = f'Epoch: [{epoch + 1}][{args.epochs}]' batch_info = f'Global Step: {global_step}' loss_info = f'Loss: {losses / args.log_step:.6f}' # Extract the learning rate from the optimizer lr = optimizer.param_groups[0]['lr'] lr_info = f'Learning Rate: {lr:.6f}' log_message = f'{current_time}\n{epoch_info} {batch_info} {loss_info} {lr_info}\n' with open(args.log_dir + 'log.txt', mode='a') as n: n.write(log_message) losses = 0.0 # check performance if (global_step + 1) % args.eval_every_step == 0: if accelerator.is_main_process: model_module = model.module if hasattr(model, 'module') else model val_psds(model_module, val_loader, params, epoch=global_step+1, split='val', save_path=args.log_dir + 'output/', device=accelerator.device) # save model unwrapped_model = accelerator.unwrap_model(model) accelerator.save({ "model": model.state_dict(), }, args.save_dir + str(global_step+1) + '.pt') accelerator.wait_for_everyone() model.train()