def main(opt):

    device = torch.device(opt.device if torch.cuda.is_available() else "cpu")
    train_set = dataset.SBU_Dataset(opt, training=True)
    test_set = dataset.SBU_Dataset(opt, training=False)
    lens = train_set.__len__()
    iters_per_epoch = math.ceil(lens / opt.batch_size)
    max_epoch = math.ceil(opt.max_iter / iters_per_epoch)

    train_loader = DataLoader(train_set, batch_size=opt.batch_size, shuffle=True,
                             num_workers=opt.job, pin_memory=True,
                             collate_fn=dataset.collate_fn, drop_last=False)
    test_loader = DataLoader(test_set, batch_size=opt.batch_size, shuffle=True,
                              num_workers=opt.job, pin_memory=True,
                              collate_fn=dataset.collate_fn, drop_last=False)

    writer = SummaryWriter()
    print("loading the model.......")
    net = VRNN(opt)
    net.to(device)
    optimizer = torch.optim.Adam(net.parameters(), lr=opt.lr, betas=(0.9, 0.999))
    best_loss = 10000
    bar = tqdm(range(max_epoch))
    for epoch in bar:
        bar.set_description('train epoch %06d' % epoch)
        train(train_loader, net, device, optimizer, writer, epoch)
        test_loss = test(test_loader, net, device, writer, epoch)
        if test_loss < best_loss:
            best_loss = test_loss
            save_model(net, optimizer,epoch)

    writer.close()
Exemple #2
0
def train(conf):

    train_loader, test_loader = load_dataset(512)
    net = VRNN(conf.x_dim, conf.h_dim, conf.z_dim)
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    torch.cuda.manual_seed_all(112858)
    net.to(device)
    net = torch.nn.DataParallel(net, device_ids=[0, 1])
    if conf.restore == True:
        net.load_state_dict(
            torch.load(conf.checkpoint_path, map_location='cuda:0'))
        print('Restore model from ' + conf.checkpoint_path)
    optimizer = optim.Adam(net.parameters(), lr=0.001)
    for ep in range(1, conf.train_epoch + 1):
        prog = Progbar(target=117)
        print("At epoch:{}".format(str(ep)))
        for i, (data, target) in enumerate(train_loader):
            data = data.squeeze(1)
            data = (data / 255).to(device)
            package = net(data)
            loss = Loss(package, data)
            net.zero_grad()
            loss.backward()
            _ = torch.nn.utils.clip_grad_norm_(net.parameters(), 5)
            optimizer.step()
            prog.update(i, exact=[("Training Loss", loss.item())])

        with torch.no_grad():
            x_decoded = net.module.sampling(conf.x_dim, device)
            x_decoded = x_decoded.cpu().numpy()
            digit = x_decoded.reshape(conf.x_dim, conf.x_dim)
            plt.imshow(digit, cmap='Greys_r')
            plt.pause(1e-6)

        if ep % conf.save_every == 0:
            torch.save(net.state_dict(),
                       '../checkpoint/Epoch_' + str(ep + 1) + '.pth')
def main(runsss, overlap, window_size, h_dim, z_dim, batch_size, language):
    try:

        def train(epoch):
            train_loss = 0
            tq = tqdm(train_loader)
            for batch_idx, (data, _) in enumerate(tq):
                data = Variable(data.squeeze().transpose(0, 1))
                data = (data - data.min().item()) / (data.max().item() -
                                                     data.min().item())
                #forward + backward + optimize
                optimizer.zero_grad()
                kld_loss, nll_loss, _, _ = model(data)
                loss = kld_loss + nll_loss
                loss.backward()
                optimizer.step()

                #grad norm clipping, only in pytorch version >= 1.10
                nn.utils.clip_grad_norm(model.parameters(), clip)

                tq.set_postfix(kld_loss=(kld_loss.item() / batch_size),
                               nll_loss=(nll_loss.item() / batch_size))
                train_loss += loss.item()
            return

        def test(epoch):
            """uses test data to evaluate 
			likelihood of the model"""

            mean_kld_loss, mean_nll_loss = 0, 0
            tq = tqdm(test_loader)
            for i, (data, _) in enumerate(tq):

                #data = Variable(data)
                data = Variable(data.squeeze().transpose(0, 1))
                data = (data - data.min().item()) / (data.max().item() -
                                                     data.min().item())

                kld_loss, nll_loss, _, _ = model(data)

                mean_kld_loss += kld_loss.item()
                mean_nll_loss += nll_loss.item()

            mean_kld_loss /= len(test_loader.dataset)
            mean_nll_loss /= len(test_loader.dataset)

            print('====> Test set loss: KLD Loss = {:.4f}, NLL Loss = {:.4f} '.
                  format(mean_kld_loss, mean_nll_loss))
            return

        def train_classifier(param_string, train_loader_enc, test_loader_enc,
                             optimizer2, classify, criterion):
            num_epochs = 100
            best = 0
            acc = []
            train_acc = []
            for epoch in range(num_epochs):
                print(f'epoch num: {epoch}\n')
                running_loss = 0.0
                tq = tqdm(train_loader_enc)
                for i, (data, labels) in enumerate(tq):

                    # zero the parameter gradients
                    optimizer2.zero_grad()

                    # forward + backward + optimize
                    outputs = classify(data)
                    # print(outputs, labels)
                    loss = criterion(outputs.float(), labels.long())
                    loss.backward()
                    optimizer2.step()
                    # print statistics
                    running_loss += loss.item()
                    tq.set_postfix(running_loss=(running_loss))
                acc.append(
                    test_classifier(train_loader_enc,
                                    test_loader_enc,
                                    classify,
                                    flag=False,
                                    param_string=param_string))
                train_acc.append(
                    test_classifier(train_loader_enc,
                                    test_loader_enc,
                                    classify,
                                    flag=True,
                                    param_string=param_string))
                print(acc[-1], best)
                if acc[-1] > best:
                    best = acc[-1]
            print(f'best acc of {best}')
            f = open('class_acc.txt', 'a+')
            for i in range(len(train_acc)):
                f.write(param_string)
                f.write(f'{i},{train_acc[i]},{acc[i]}\n')
            f.close()
            return max([*acc, *train_acc])

        def test_classifier(train_loader_enc,
                            test_loader_enc,
                            classify,
                            flag=False,
                            param_string=""):
            # evaluate
            correct = 0
            total = 0
            classes = defaultdict(int)
            wrong = defaultdict(int)
            y_pred = []
            y_true = []
            with torch.no_grad():
                if flag:
                    tq = tqdm(train_loader_enc)
                    for data in tq:
                        dat, labels = data
                        outputs = classify(dat)
                        for i in range(len(outputs)):
                            if outputs[i][0] >= outputs[i][1]:
                                pred = 0
                            else:
                                pred = 1
                            y_pred.append(pred)
                            y_true.append(labels[i].item())
                            if pred == labels[i].item():
                                correct += 1
                                classes[pred] += 1
                            else:
                                wrong[pred] += 1
                            total += 1
                else:
                    tq = tqdm(test_loader_enc)
                    for data in tq:
                        dat, labels = data
                        outputs = classify(dat)
                        for i in range(len(outputs)):
                            if outputs[i][0] >= outputs[i][1]:
                                pred = 0
                            else:
                                pred = 1
                            y_pred.append(pred)
                            y_true.append(labels[i].item())
                            if pred == labels[i].item():
                                correct += 1
                                classes[pred] += 1
                            else:
                                wrong[pred] += 1
                            total += 1
            f11 = f1_score(y_true, y_pred, average='binary')
            [precision, recall, fbeta_score,
             support] = precision_recall_fscore_support(y_true,
                                                        y_pred,
                                                        average='binary')
            paramms = f'correct: {correct}, total:  {total}, classes:  {classes}, wrong: {wrong}, f1_score:  {f11}, precision: {precision}, recall: {recall}, fbeta_score: {fbeta_score}, support: {support}'
            print(paramms)

            print(f'accuracy: {correct/total}')
            if param_string:
                f = open('class_acc.txt', 'a+')
                f.write(param_string)
                f.write(f'y_pred-{y_pred},y_true-{y_true}\n')
                f.write(f'{paramms}\n')
                f.close()
            acc = correct / total
            return acc

        # transform inputs from test set to encoded vectors, make new training training loaders
        def transform_inputs(loader, batch_size=batch_size):
            encoded_inputs = []
            labels = []
            tq = tqdm(loader)
            with torch.no_grad():
                for batch_idx, (data, label) in enumerate(tq):
                    data = Variable(data.squeeze().transpose(0, 1))
                    data = (data - data.min().item()) / (data.max().item() -
                                                         data.min().item())
                    h = model.predict(data)
                    for i in range(h.shape[1]):
                        encoded_inputs.append(h[:, i, :].flatten().numpy())
                        labels.append(label[i].item())
            return torch.utils.data.DataLoader(torch.utils.data.TensorDataset(
                torch.Tensor(encoded_inputs), torch.Tensor(labels)),
                                               batch_size=batch_size,
                                               shuffle=True)

        def run_classifier(epochh, fn):
            for b_size in [4, 8, 16, 32]:
                train_loader_enc = transform_inputs(train_loader, b_size)
                test_loader_enc = transform_inputs(test_loader, b_size)
                for intermediate_dim in [5, 10, 20, 40]:
                    for layers in [True, False]:
                        f = open(txt_file, 'a+')
                        f.write(
                            f'fn="{fn}",runsss={runsss},epochh={epochh},overlap={overlap},window_size={window_size},h_dim={h_dim},z_dim={z_dim},batch_size={batch_size},language={language},b_size={b_size},intermediate_dim={intermediate_dim},layers={layers}\n'
                        )
                        f.close()

                        classify = Classifier(input_dim=h_dim,
                                              intermediate_dim=20,
                                              layers=layers)
                        print(classify, classify.count_parameters())
                        criterion = nn.CrossEntropyLoss()
                        optimizer2 = torch.optim.Adam(classify.parameters(),
                                                      lr=0.001)

                        train_classifier(
                            f'fn="{fn}",runsss={runsss},epochh={epochh},overlap={overlap},window_size={window_size},h_dim={h_dim},z_dim={z_dim},batch_size={batch_size},language={language},b_size={b_size},intermediate_dim={intermediate_dim},layers={layers}\n',
                            train_loader_enc, test_loader_enc, optimizer2,
                            classify, criterion)
                        accuracy = test_classifier(train_loader_enc,
                                                   test_loader_enc,
                                                   classify,
                                                   flag=False)

                        print(f'final accuracy ---> {accuracy}')
                        print('Finished Training')
            return

        x = get_dataset2(overlap=overlap,
                         window_size=window_size,
                         time_steps=20,
                         language=language,
                         max_len=(100 if language == "english" else 80))
        #hyperparameters
        # x_dim = 70
        x_dim = x['genuine'].shape[2]
        # h_dim = 100
        # z_dim = 16
        n_layers = 1
        n_epochs = 25
        clip = 10
        learning_rate = 3e-4
        batch_size = batch_size
        seed = 128
        print_every = 10
        save_every = 5

        #manual seed
        torch.manual_seed(seed)

        #init model + optimizer + datasets
        x_i = []
        y_i = []
        x_it = []
        y_it = []
        test_count = defaultdict(int)
        tot = 80
        for val in ['genuine', 'forged']:
            for i in x[val]:
                if val == 'genuine' and test_count['genuine'] < tot:
                    x_it.append(i)
                    y_it.append([1])
                    test_count['genuine'] += 1
                elif val == 'forged' and test_count['forged'] < tot:
                    x_it.append(i)
                    y_it.append([0])
                    test_count['forged'] += 1
                else:
                    x_i.append(i)
                    y_i.append([0] if val == 'genuine' else [1])

        # print(len(x_i),len(y_i),len(x_it),len(y_it))
        x_i, y_i, x_it, y_it = np.array(x_i), np.array(y_i).reshape(
            (-1, )), np.array(x_it), np.array(y_it).reshape((-1, ))

        if False:
            signatures_train, signatures_test, labels_train, labels_test = get_dataset(
            )
        else:
            signatures_train, signatures_test, labels_train, labels_test = x_i, x_it, y_i, y_it

        print('input data\n', signatures_train.shape, labels_train.shape,
              signatures_test.shape, labels_test.shape)

        x_dim = signatures_train.shape[2]

        train_loader = torch.utils.data.DataLoader(
            torch.utils.data.TensorDataset(torch.Tensor(signatures_train),
                                           torch.Tensor(labels_train)),
            batch_size=batch_size,
            shuffle=True)
        test_loader = torch.utils.data.DataLoader(
            torch.utils.data.TensorDataset(torch.Tensor(signatures_test),
                                           torch.Tensor(labels_test)),
            batch_size=batch_size,
            shuffle=True)

        model = VRNN(x_dim, h_dim, z_dim, n_layers)

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

        for epoch in range(1, n_epochs + 1):

            #training + testing
            train(epoch)
            test(epoch)

            #saving model
            if epoch % save_every == 1:
                fn = 'saves/vrnn_state_dict_' + f'{runsss},{overlap},{window_size},{h_dim},{z_dim},{batch_size}.{language},{epoch}' + '.pth'
                torch.save(model.state_dict(), fn)
                print('Saved model to ' + fn)
                run_classifier(epoch, fn + '\n')

        # freeze model weights
        for param in model.parameters():
            param.requires_grad = False

    except:
        print('FAILED RUN')
        f = open(txt_file, 'a+')
        f.write(
            f'FAILED RUN ---> {runsss},{overlap},{window_size},{h_dim},{z_dim},{batch_size}.{language}\n'
        )
        f.close()
save_every = 10

#manual seed
torch.manual_seed(seed)
plt.ion()

#init model + optimizer + datasets
train_loader = torch.utils.data.DataLoader(datasets.MNIST(
    'data', train=True, download=True, transform=transforms.ToTensor()),
                                           batch_size=batch_size,
                                           shuffle=True)

test_loader = torch.utils.data.DataLoader(datasets.MNIST(
    'data', train=False, transform=transforms.ToTensor()),
                                          batch_size=batch_size,
                                          shuffle=True)

model = VRNN(x_dim, h_dim, z_dim, n_layers)
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

for epoch in range(1, n_epochs + 1):

    #training + testing
    train(epoch)
    test(epoch)

    #saving model
    if epoch % save_every == 1:
        fn = 'saves/vrnn_state_dict_' + str(epoch) + '.pth'
        torch.save(model.state_dict(), fn)
        print('Saved model to ' + fn)
Exemple #5
0
def main(args):
    # Create model directory
    if not os.path.exists(args.model_path):
        os.makedirs(args.model_path)

    # Load vocabulary wrapper.
    with open(args.vocab_path, 'rb') as f:
        vocab = pickle.load(f)

    # Image preprocessing
    # For normalization, see https://github.com/pytorch/vision#models
    transform = transforms.Compose([
        transforms.RandomCrop(args.crop_size),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
    ])
    #val_loader = get_loader('./data/val_resized2014/', './data/annotations/captions_val2014.json',
    #                         vocab, transform, 1, False, 1)

    start_epoch = 0

    encoder_state = args.encoder
    decoder_state = args.decoder

    # Build the models
    encoder = EncoderCNN(args.embed_size)
    if not args.train_encoder:
        encoder.eval()
    decoder = VRNN(args.embed_size, args.hidden_size, len(vocab),
                   args.latent_size, args.num_layers)

    if args.restart:
        encoder_state, decoder_state = 'new', 'new'

    if encoder_state == '': encoder_state = 'new'
    if decoder_state == '': decoder_state = 'new'

    print("Using encoder: {}".format(encoder_state))
    print("Using decoder: {}".format(decoder_state))

    try:
        start_epoch = int(float(decoder_state.split('-')[1]))
    except:
        pass

    if encoder_state != 'new':
        encoder.load_state_dict(torch.load(encoder_state))
    if decoder_state != 'new':
        decoder.load_state_dict(torch.load(decoder_state))

    # Build data loader
    data_loader = get_loader(args.image_dir,
                             args.caption_path,
                             vocab,
                             transform,
                             args.batch_size,
                             shuffle=True,
                             num_workers=args.num_workers)
    """ Make logfile and log output """
    with open(args.model_path + args.logfile, 'a+') as f:
        f.write("Using encoder: new\nUsing decoder: new\n\n")

    if torch.cuda.is_available():
        encoder.cuda()
        decoder.cuda()

    # Optimizer
    cross_entropy = nn.CrossEntropyLoss()
    params = list(decoder.parameters()) + list(
        encoder.linear.parameters()) + list(encoder.bn.parameters())
    optimizer = torch.optim.Adam(params, lr=args.learning_rate)

    batch_loss = []
    batch_loss_det = []
    batch_kl = []
    batch_ml = []
    batch_acc = []

    # Train the Models
    total_step = len(data_loader)
    for epoch in range(start_epoch, args.num_epochs):
        for i, (images, captions, lengths, _, _) in enumerate(data_loader):

            # get lengths excluding <start> symbol
            lengths = [l - 1 for l in lengths]

            # Set mini-batch dataset
            images = to_var(images, volatile=True)
            captions = to_var(captions)

            # assuming following assertion
            assert min(lengths) > args.z_step + 2

            # get targets from captions (excluding <start> tokens)
            #targets = pack_padded_sequence(captions[:,1:], lengths, batch_first=True)[0]
            targets_var = captions[:, args.z_step + 1]
            targets_det = pack_padded_sequence(
                captions[:, args.z_step + 2:],
                [l - args.z_step - 1 for l in lengths],
                batch_first=True)[0]

            # Get prior and approximate distributions
            decoder.zero_grad()
            encoder.zero_grad()
            features = encoder(images)
            prior, q_z, q_x, det_x = decoder(features,
                                             captions,
                                             lengths,
                                             z_step=args.z_step)

            # Calculate KL Divergence
            kl = torch.mean(kl_divergence(*q_z + prior))

            # Get marginal likelihood from log likelihood of the correct symbol
            index = (torch.cuda.LongTensor(range(q_x.shape[0])), targets_var)
            ml = torch.mean(q_x[index])

            # Get Cross-Entropy loss for deterministic decoder
            ce = cross_entropy(det_x, targets_det)

            elbo = ml - kl
            loss_var = -elbo

            loss_det = ce

            loss = loss_var + loss_det

            batch_loss.append(loss.data[0])
            batch_loss_det.append(loss_det.data[0])
            batch_kl.append(kl.data[0])
            batch_ml.append(ml.data[0])

            loss.backward()
            optimizer.step()

            # Print log info
            if i % args.log_step == 0:
                print(
                    'Epoch [%d/%d], Step [%d/%d], Loss: %.4f, Perplexity: %5.4f'
                    % (epoch, args.num_epochs, i, total_step, loss.data[0],
                       np.exp(loss.data[0])))

                with open(args.model_path + args.logfile, 'a') as f:
                    f.write(
                        'Epoch [%d/%d], Step [%d/%d], Loss: %.4f, Perplexity: %5.4f'
                        % (epoch, args.num_epochs, i, total_step, loss.data[0],
                           np.exp(loss.data[0])))

            # Save the models
            if (i + 1) % args.save_step == 0:
                torch.save(
                    decoder.state_dict(),
                    os.path.join(args.model_path,
                                 'decoder-%d-%d.pkl' % (epoch + 1, i + 1)))
                if args.train_encoder:
                    torch.save(
                        encoder.state_dict(),
                        os.path.join(args.model_path,
                                     'encoder-%d-%d.pkl' % (epoch + 1, i + 1)))
                with open(args.model_path + 'training_loss.pkl', 'w+') as f:
                    pickle.dump(batch_loss, f)
                with open(args.model_path + 'training_val.pkl', 'w+') as f:
                    pickle.dump(batch_acc, f)

    with open(args.model_path + args.logfile, 'a') as f:
        f.write("Training finished at {} .\n\n".format(str(datetime.now())))