from load_mnist import load_mnist_classSelect
X, Y, tridx = load_mnist_classSelect('train', data_classes, ylabels)
vaX, vaY, vaidx = load_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
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'])
# --- train/load GCE ---
from models.CVAE import Decoder, Encoder
if retrain_gce:
encoder = Encoder(K + L, c_dim, x_dim).to(device)
decoder = Decoder(K + L, c_dim, x_dim).to(device)
encoder.apply(util.weights_init_normal)
decoder.apply(util.weights_init_normal)
gce = GenerativeCausalExplainer(classifier, decoder, encoder, device)
traininfo = gce.train(X,
K,
L,
steps=train_steps,
Nalpha=Nalpha,
Nbeta=Nbeta,
lam=lam,
batch_size=batch_size,
lr=lr)
if save_gce:
if not os.path.exists(gce_path):
# --- initialize VAE and train GCE ---
from models.CVAE import Decoder, Encoder
data = {
'loss': np.zeros((len(filts_per_layer), len(lambdas), train_steps)),
'loss_ce': np.zeros((len(filts_per_layer), len(lambdas), train_steps)),
'loss_nll': np.zeros((len(filts_per_layer), len(lambdas), train_steps)),
'Ijoint': np.zeros((len(filts_per_layer), len(lambdas))),
'Is': np.zeros((len(filts_per_layer), len(lambdas), K + L))
}
for (i_f, nfilt) in enumerate(filts_per_layer):
for (i_l, lam) in enumerate(lambdas):
filename = 'model_%dfilters_lambda%g.pt' % (nfilt, lam)
if retrain_gce:
print('=== %d FILTERS PER LAYER, LAMBDA = %g ===' % (nfilt, lam))
# initialize VAE
encoder = Encoder(K + L, c_dim, x_dim,
filt_per_layer=nfilt).to(device)
decoder = Decoder(K + L, c_dim, x_dim,
filt_per_layer=nfilt).to(device)
encoder.apply(util.weights_init_normal)
decoder.apply(util.weights_init_normal)
# train GCE
gce = GenerativeCausalExplainer(classifier,
decoder,
encoder,
device,
debug_print=False)
traininfo = gce.train(X,
K,
L,
steps=train_steps,
Nalpha=Nalpha,