How to use the torch.save function in torch

To help you get started, we’ve selected a few torch examples, based on popular ways it is used in public projects.

Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.

github yxlijun / S3FD.pytorch / train.py View on Github external
loss_l, loss_c = criterion(out, targets)
        loss = loss_l + loss_c
        loc_loss += loss_l.data[0]
        conf_loss += loss_c.data[0]
        step += 1

    tloss = (loc_loss + conf_loss) / step
    t2 = time.time()
    print('Timer: %.4f' % (t2 - t1))
    print('test epoch:' + repr(epoch) + ' || Loss:%.4f' % (tloss))

    global min_loss
    if tloss < min_loss:
        print('Saving best state,epoch', epoch)
        file = 'sfd_{}.pth'.format(args.dataset)
        torch.save(s3fd_net.state_dict(), os.path.join(
            args.save_folder, file))
        min_loss = tloss

    states = {
        'epoch': epoch,
        'weight': s3fd_net.state_dict(),
    }
    file = 'sfd_{}_checkpoint.pth'.format(args.dataset)
    torch.save(states, os.path.join(
        args.save_folder, file))
github KaihuaTang / VCTree-Visual-Question-Answering / train.py View on Github external
'eval': {
                    'answers': r[0],
                    'accuracies': r[1],
                    'idx': r[2],
                },
                'vocab': val_loader.dataset.vocab,
                'src': src,
                'setting': exp_setting,
            }
            current_ac = sum(r[1]) / len(r[1])
            if current_ac >  best_accuracy:
                best_accuracy = current_ac
                print('update best model, current: ', current_ac)
                torch.save(results, target_name + '_best.pth')
            if i % 1 == 0:
                torch.save(results, target_name + '_' + str(i) + '.pth')

        else:
            # in test mode, save a results file in the format accepted by the submission server
            answer_index_to_string = {a:  s for s, a in val_loader.dataset.answer_to_index.items()}
            results = []
            for answer, index in zip(r[0], r[2]):
                answer = answer_index_to_string[answer.item()]
                qid = val_loader.dataset.question_ids[index]
                entry = {
                    'question_id': qid,
                    'answer': answer,
                }
                results.append(entry)
            with open('results.json', 'w') as fd:
                json.dump(results, fd)
github carpedm20 / BEGAN-pytorch / trainer.py View on Github external
def save_model(self, step):
        print("[*] Save models to {}...".format(self.model_dir))

        torch.save(self.G.state_dict(), '{}/G_{}.pth'.format(self.model_dir, step))
        torch.save(self.D.state_dict(), '{}/D_{}.pth'.format(self.model_dir, step))
github hongweizeng / cnn-seq2seq / train.py View on Github external
#  (3) update the learning rate
        optim.updateLearningRate(valid_loss, epoch)

        model_state_dict = model.module.state_dict() if len(opt.gpus) > 1 else model.state_dict()
        model_state_dict = {k: v for k, v in model_state_dict.items() if 'generator' not in k}
        generator_state_dict = model.generator.module.state_dict() if len(opt.gpus) > 1 else model.generator.state_dict()
        #  (4) drop a checkpoint
        checkpoint = {
            'model': model_state_dict,
            'generator': generator_state_dict,
            'dicts': dataset['dicts'],
            'opt': opt,
            'epoch': epoch,
            'optim': optim
        }
        torch.save(checkpoint,
                   '%s_acc_%.2f_ppl_%.2f_e%d.pt' % (opt.save_model, 100*valid_acc, valid_ppl, epoch))
github XuezheMax / macow / experiments / celeba.py View on Github external
# number of parameters
print('# of Parameters: %d' % (sum([param.numel() for param in fgen.parameters()])))
fgen.sync()
lr_min = lr / 100
lr = scheduler.get_lr()[0]
for epoch in range(start_epoch, args.epochs + 1):
    train(epoch, train_k)
    fgen.sync()
    print('-' * 100)
    with torch.no_grad():
        nll_mc, nent, nll_iw, bpd_mc, nepd, bpd_iw = eval(test_loader, test_k)

    if nll_mc < best_nll_mc:
        patient = 0
        torch.save(fgen.state_dict(), model_name)

        best_epoch = epoch
        best_nll_mc = nll_mc
        best_nll_iw = nll_iw
        best_bpd_mc = bpd_mc
        best_bpd_iw = bpd_iw
        best_nent = nent
        best_nepd = nepd

        with torch.no_grad():
            try:
                reconstruct(epoch)
            except RuntimeError:
                print('reconstruction failed.')
            try:
                sample(epoch)
github majumderb / recipe-personalization / recipe_gen / analysis / next_step_pred.py View on Github external
tr_loss += loss.item()
                nb_tr_examples += input_ids.size(0)
                nb_tr_steps += 1
                if (step + 1) % args.gradient_accumulation_steps == 0:
                    if args.fp16:
                        # modify learning rate with special warm up BERT uses
                        # if args.fp16 is False, BertAdam is used that handles this automatically
                        lr_this_step = args.learning_rate * warmup_linear(global_step/num_train_optimization_steps, args.warmup_proportion)
                        for param_group in optimizer.param_groups:
                            param_group['lr'] = lr_this_step
                    optimizer.step()
                    optimizer.zero_grad()
                    global_step += 1
    
    # save model
    torch.save(model.state_dict(), os.path.join(args.output_dir, 'nsp_model.pt'))

    if args.do_eval and (args.local_rank == -1 or torch.distributed.get_rank() == 0):
        eval_examples = processor.get_dev_examples(args.data_dir)
        eval_features = convert_examples_to_features(
            eval_examples, label_list, args.max_seq_length, tokenizer)
        logger.info("***** Running evaluation *****")
        logger.info("  Num examples = %d", len(eval_examples))
        logger.info("  Batch size = %d", args.eval_batch_size)
        all_input_ids = torch.tensor([f.input_ids for f in eval_features], dtype=torch.long)
        all_input_mask = torch.tensor([f.input_mask for f in eval_features], dtype=torch.long)
        all_segment_ids = torch.tensor([f.segment_ids for f in eval_features], dtype=torch.long)
        all_label_ids = torch.tensor([f.label_id for f in eval_features], dtype=torch.long)
        eval_data = TensorDataset(all_input_ids, all_input_mask, all_segment_ids, all_label_ids)
        # Run prediction for full data
        eval_sampler = SequentialSampler(eval_data)
        eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
github SforAiDl / Neural-Voice-Cloning-With-Few-Samples / train_whole.py View on Github external
def save_checkpoint(model, optimizer, checkpoint_path, epoch):

    optimizer_state = optimizer.state_dict()
    torch.save({
        "state_dict": model.state_dict(),
        "optimizer": optimizer_state,
        "global_epoch": epoch,
        "epoch":epoch+1,

    }, checkpoint_path)
    print("Saved checkpoint:", checkpoint_path)
github xieydd / Pytorch-Single-Path-One-Shot-NAS / utils / utils.py View on Github external
def save_checkpoint(self, val_loss, model, ckpt_dir,is_best=False):
        '''Saves model when validation loss decrease.'''
        if self.verbose:
            print('Validation loss decreased ({:.6f} --> {:.6f}).  Saving model ...'.format(self.val_loss_min, val_loss))
        filename = os.path.join(ckpt_dir, 'checkpoint.pth.tar')
        torch.save(model, filename)
        if is_best:
            best_filename = os.path.join(ckpt_dir, 'best.pth.tar')
            shutil.copyfile(filename, best_filename)
        self.val_loss_min = val_loss
github sebgao / LIP / imagenet / main.py View on Github external
def save_checkpoint(state, is_best, args, filename='checkpoint.pth.tar'):
    torch.save(state, os.path.join(args.work_dirs, filename))
    if is_best:
        shutil.copyfile(os.path.join(args.work_dirs, filename), os.path.join(args.work_dirs, 'model_best.pth.tar'))
github HelloRicky123 / Siamese-RPN / 54version / siamfc / train.py View on Github external
pred_box = box_transform_inv(anchors, topk_offset)
                        img_box = add_box_img(inst_img, pred_box)
                        cls_pred = conf_target[0]
                        gt_box = get_topk_box(cls_pred, regression_target[0], anchors_show)
                        img_box = add_box_img(img_box, gt_box, color=(255, 0, 0))
                        vis.plot_img(img_box.transpose(2, 0, 1), win=6, name='box_max_iou')
        adjust_learning_rate(optimizer, 1 / config.warm_scale)

        save_name = "./models/siamrpn_warm.pth"
        new_state_dict = model.state_dict()
        if torch.cuda.device_count() > 1:
            new_state_dict = OrderedDict()
            for k, v in model.state_dict().items():
                namekey = k[7:]  # remove `module.`
                new_state_dict[namekey] = v
        torch.save({
            'epoch': 0,
            'model': new_state_dict,
            'optimizer': optimizer.state_dict(),
        }, save_name)
        print('save model: {}'.format(save_name))

    for epoch in range(start_epoch, config.epoch + 1):
        train_loss = []
        model.train()
        loss_temp_cls = 0
        loss_temp_reg = 0
        for i, data in enumerate(tqdm(trainloader)):
            exemplar_imgs, instance_imgs, regression_target, conf_target = data
            # conf_target (8,1125) (8,225x5)
            regression_target, conf_target = regression_target.cuda(), conf_target.cuda()