def load_model(use_cuda, log_dir, cp_num, embedding_size, n_classes): model = background_resnet(embedding_size=embedding_size, num_classes=n_classes) if use_cuda: model.cuda() print('=> loading checkpoint') # original saved file with DataParallel checkpoint = torch.load(log_dir + '/checkpoint_' + str(cp_num) + '.pth') # create new OrderedDict that does not contain `module.` model.load_state_dict(checkpoint['state_dict']) model.eval() return model
def load_model(use_cuda, log_dir, cp_num, n_classes): model = background_resnet(num_classes=n_classes) if use_cuda: model.cuda() print('=> loading checkpoint') # load pre-trained parameters checkpoint = torch.load(log_dir + '/checkpoint_' + str(cp_num).zfill(3) + '.pth') model.load_state_dict(checkpoint['state_dict']) model.eval() return model
def load_model(log_dir, cp_num, n_classes): model = background_resnet(num_classes=n_classes) optimizer = create_optimizer(model) print('=> loading checkpoint') checkpoint = torch.load(log_dir + '/checkpoint_' + str(cp_num).zfill(3) + '.pth') # create new OrderedDict that does not contain `module.` model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) for state in optimizer.state.values(): for k, v in state.items(): if isinstance(v, torch.Tensor): state[k] = v.cuda() return model, optimizer
def main(): if args.seed is not None: random.seed(args.seed) torch.manual_seed(args.seed) np.random.seed(args.seed) cudnn.deterministic = True if args.use_cuda: torch.cuda.manual_seed(args.seed) torch.cuda.manual_seed_all(args.seed) warnings.warn('You have chosen to seed training. ' 'This will turn on the CUDNN deterministic setting, ' 'which can slow down your training considerably! ' 'You may see unexpected behavior when restarting ' 'from checkpoints.') if args.use_checkpoint: start = args.cp_num + 1 else: start = 0 # Start epoch n_epochs = max_epoch - start # How many epochs? # Load dataset train_DB, n_data, n_classes = make_DB(DB_type=args.data_type) n_episode = int(n_data / ((args.n_shot + args.n_query) * args.nb_class_train)) if not os.path.exists(log_dir): os.makedirs(log_dir) # Generate model and optimizer if args.use_checkpoint: model, optimizer = load_model(log_dir, args.cp_num, n_classes) else: model = background_resnet(num_classes=n_classes) optimizer = create_optimizer(model) # define objective function, optimizer and scheduler objective = Prototypical() if args.loss_type == 'prototypical' else SoftmaxLoss() scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience=3, min_lr=1e-5, threshold=1e-4, verbose=1) if args.use_cuda: model.cuda() train_generator = metaGenerator(train_DB, read_MFB, nb_classes=args.nb_class_train, nb_samples_per_class=args.n_shot + args.n_query, max_iter=n_episode * (n_epochs-args.cp_num), xp=np) # training train(train_generator, model, objective, optimizer, n_episode, log_dir, scheduler)
def main(): # Set hyperparameters use_cuda = True # use gpu or cpu val_ratio = 10 # Percentage of validation set embedding_size = 128 start = 1 # Start epoch n_epochs = 30 # How many epochs? end = start + n_epochs # Last epoch lr = 1e-1 # Initial learning rate wd = 1e-4 # Weight decay (L2 penalty) optimizer_type = 'sgd' # ex) sgd, adam, adagrad batch_size = 64 # Batch size for training valid_batch_size = 16 # Batch size for validation use_shuffle = True # Shuffle for training or not # Load dataset train_dataset, valid_dataset, n_classes = load_dataset(val_ratio) # print the experiment configuration print('\nNumber of classes (speakers):\n{}\n'.format(n_classes)) log_dir = 'model_saved' # where to save checkpoints if not os.path.exists(log_dir): os.makedirs(log_dir) # instantiate model and initialize weights model = background_resnet(embedding_size=embedding_size, num_classes=n_classes) if use_cuda: model.cuda() # define loss function (criterion), optimizer and scheduler criterion = nn.CrossEntropyLoss() optimizer = create_optimizer(optimizer_type, model, lr, wd) scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience=2, min_lr=1e-4, verbose=1) train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=use_shuffle) valid_loader = torch.utils.data.DataLoader(dataset=valid_dataset, batch_size=valid_batch_size, shuffle=False, collate_fn = collate_fn_feat_padded) # to track the average training loss per epoch as the model trains avg_train_losses = [] # to track the average validation loss per epoch as the model trains avg_valid_losses = [] for epoch in range(start, end): # train for one epoch train_loss = train(train_loader, model, criterion, optimizer, use_cuda, epoch, n_classes) # evaluate on validation set valid_loss = validate(valid_loader, model, criterion, use_cuda, epoch) scheduler.step(valid_loss, epoch) # calculate average loss over an epoch avg_train_losses.append(train_loss) avg_valid_losses.append(valid_loss) # do checkpointing torch.save({'epoch': epoch + 1, 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict()}, '{}/checkpoint_{}.pth'.format(log_dir, epoch)) # find position of lowest validation loss minposs = avg_valid_losses.index(min(avg_valid_losses))+1 print('Lowest validation loss at epoch %d' %minposs) # visualize the loss and learning rate as the network trained visualize_the_losses(avg_train_losses, avg_valid_losses)
def main(): # GPU parmas idx_cuda = 2 if idx_cuda < 0: device = torch.device('cpu') else: device = torch.device(f'cuda:{idx_cuda}') # Set hyperparameters embedding_size = 128 # origial 128 start = 1 # Start epoch n_epochs = 60 # How many epochs? end = start + n_epochs # Last epoch lr = 1e-1 # Initial learning rate wd = 1e-4 # Weight decay (L2 penalty) optimizer_type = 'adam' # ex) sgd, adam, adagrad # The effective batch size is M*N M = 5 # number of utterances per speaker N = 10 # Number of speaker print(f'Number of utternaces per speaker: {M}') print(f'Number of speakers: {N}') print(f'** Total batch size: {M*N} **') use_shuffle = True # Shuffle for training or not save_interval = 5 # Load dataset train_dataset, n_classes = load_dataset(M) log_dir = 'model_saved_verification_2' # where to save checkpoints if not os.path.exists(log_dir): os.makedirs(log_dir) # instantiate model and initialize weights model = background_resnet(embedding_size=embedding_size, num_classes=n_classes, backbone='resnet34') # Load the wights trained for identification #smodel.load_state_dict(torch.load('model_saved/checkpoint_50.pth')['state_dict']) # remove the last layers model.to(device) # define loss function (criterion), optimizer and scheduler criterion = GE2ELoss(device) optimizer = create_optimizer(optimizer_type, model, lr, wd) # Define dataloader train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=N, shuffle=use_shuffle, collate_fn=my_collate) # to track the average training loss per epoch as the model trains avg_train_losses = [] for epoch in range(start, end): # train for one epoch train_loss = train(train_loader, model, criterion, optimizer, device, epoch, N, M) # calculate average loss over an epoch avg_train_losses.append(train_loss) # do checkpointing if epoch % save_interval == 0 or epoch == end - 1: torch.save( { 'epoch': epoch + 1, 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict() }, '{}/checkpoint_{}.pth'.format(log_dir, epoch))