state = torch.load('./models/AE_LSUN_32_code_size_'+str(nb_classes)+'_classes.pth')
rec_model.load_state_dict(state)

print('Reconstructing data')
trainset = reconstruct_dataset_with_AE(trainset, rec_model.cuda(), bs = 1000, real_data_ratio=0)  
testset = reconstruct_dataset_with_AE(testset, rec_model.cuda(), bs = 1000, real_data_ratio=0)

train_loader = data_utils.DataLoader(trainset, batch_size=batch_size, shuffle = True)
test_loader = data_utils.DataLoader(testset, batch_size=batch_size, shuffle = False)
#rec_model = autoencoder(32)
#state = torch.load('./models/AE_32_code_size_500_classes_2000_samples.pth')
#rec_model.load_state(state)
pretrained_model_dict = torch.load('../pretrained_models/batch_classifier_LSUN_30_classes.pth')
pretrained_model = new_models.Classifier_2048_features(30)
pretrained_model.load_state_dict(pretrained_model_dict)
acc = test_model(pretrained_model.cuda(), train_loader)
print('Accuracy of pretrained model on erconstructed dataset: ' + str(acc))
print('Training')

model = Net(nb_classes).cuda()
criterion = nn.CrossEntropyLoss().cuda()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.99)
data_size = len(trainset[0])
max_test_acc = 0
for epoch in range(epochs):  # loop over the dataset multiple times
    running_loss = 0.0
#    for idx, data in enumerate(trainloader, 0):
    for idx, (train_X, train_Y) in enumerate(train_loader):
      inputs = train_X.cuda()
      labels = train_Y.cuda()
        # zero the parameter gradients
        optimizer.zero_grad()

        # forward + backward + optimize
        outputs = model(inputs)
        loss = criterion(outputs, labels.long())
        loss.backward()
        optimizer.step()

        # print statistics
        running_loss += loss.item()
        if idx % 50 == 49:  # print every 2000 mini-batches
            print('[%d, %5d] loss: %.3f' %
                  (epoch + 1, idx + 1, running_loss / 100))
            running_loss = 0.0

    test_acc = test_model(model, test_loader)
    if test_acc > max_test_acc:
        max_test_acc = test_acc
        best_model = model.float()
        #torch.save(best_model, './results/synthetic/models/'+dataset+'_classifier_original.pth')

    print('Test accuracy: ' + str(test_acc))

torch.save(
    model, './models/batch_classifier_' + str(nb_classes) + '_classes_' +
    str(train_class_size) + '_samples.pth')
#torch.save(data_sampler, './models/data_sampler_'+ str(nb_classes) +'_classes.pth')
print('Test accuracy')
print('Finished Training')
classifier_model_dict = torch.load(
    '../pretrained_models/batch_classifier_LSUN_30_classes.pth')
classifier_model = new_models.Classifier_2048_features(30)
classifier_model.load_state_dict(classifier_model_dict)
classifier_model = classifier_model.cuda()
criterion_AE = nn.MSELoss().cuda()
criterion_classif = nn.MSELoss().cuda()
#optimizer = torch.optim.SGD(autoencoder_model.parameters(), lr=opts['learning_rate'], momentum=0.99)
optimizer_main = torch.optim.Adam(autoencoder_model.parameters(),
                                  lr=opts['learning_rate'],
                                  betas=(0.9, 0.999),
                                  weight_decay=1e-5)

accuracies = []
best_acc = 0
acc = test_model(classifier_model, test_loader)
print('Accuracy of pretrained model on the original testset: ' + str(acc))
for epoch in range(opts['number_of_epochs']):
    bar = Bar('Training: ',
              max=int(opts['nb_classes'] * 100000 / opts['batch_size']))
    for idx, (train_X, train_Y) in enumerate(train_loader):
        bar.next()
        inputs = train_X.cuda()
        labels = train_Y.cuda()
        optimizer_main.zero_grad()
        #optimizer_class.zero_grad()
        #
        #    img = Variable(inputs).cuda()
        # ===================forward=====================
        outputs = autoencoder_model(inputs)