lambdas = np.logspace(-3, -1, 10)
# other
randseed = 0
gce_path = './pretrained_models/vae_capacity'
retrain_gce = True
save_gce = True
# --- initialize ---
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
if randseed is not None:
np.random.seed(randseed)
torch.manual_seed(randseed)
# --- load data ---
from load_mnist import load_fashion_mnist_classSelect
X, Y, tridx = load_fashion_mnist_classSelect('train', data_classes,
range(0, len(data_classes)))
vaX, vaY, vaidx = load_fashion_mnist_classSelect('val', data_classes,
range(0, len(data_classes)))
ntrain, nrow, ncol, c_dim = X.shape
x_dim = nrow * ncol
# --- load classifier ---
from models.CNN_classifier import CNN
classifier = CNN(len(data_classes)).to(device)
checkpoint = torch.load('%s/model.pt' % classifier_path, map_location=device)
classifier.load_state_dict(checkpoint['model_state_dict_classifier'])
# --- initialize VAE and train GCE ---
from models.CVAE import Decoder, Encoder
data = {
'loss': np.zeros((len(filts_per_layer), len(lambdas), train_steps)),
z_dim = K + L
y_dim = data_classes.shape[0]
ylabels = range(0, y_dim)
nsweep = len(latent_sweep_vals)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
params = {
'z_dim': K + L,
'alpha_dim': K,
'No': 100,
'Ni': 25,
'decoder_net': 'VAE_CNN',
'break_up_ce': False
}
# --- load test data ---
X, Y, tridx = load_mnist.load_fashion_mnist_classSelect(
'train', data_classes, ylabels)
vaX, vaY, vaidx = load_mnist.load_fashion_mnist_classSelect(
'val', data_classes, ylabels)
ntrain, nrow, ncol, c_dim = X.shape
x_dim = nrow * ncol
# --- load classifier ---
from models.CNN_classifier import CNN
checkpoint_model = torch.load(os.path.join(classifier_path, 'model.pt'),
map_location=device)
classifier = CNN(y_dim).to(device)
classifier.load_state_dict(checkpoint_model['model_state_dict_classifier'])
# --- train/load GCE ---
from models.CVAE import Decoder, Encoder
if retrain_gce: