def check_sgd_gmm(D, K, N, plot=False, device=None):
if not device:
device = torch.device('cpu')
means = (np.random.rand(K, D) * 20) - 10
q = (2 * np.random.randn(K, D, D))
covars = np.matmul(q.swapaxes(1, 2), q)
X = np.empty((N, K, D))
for i in range(K):
X[:, i, :] = np.random.multivariate_normal(
mean=means[i, :],
cov=covars[i, :, :],
size=N
)
X_data = [torch.Tensor(X.reshape(-1, D).astype(np.float32))]
gmm = SGDGMM(K, D, device=device, epochs=200)
gmm.fit(X_data)
if plot:
fig, ax = plt.subplots()
for i in range(K):
sc = ax.scatter(
X[:, i, 0],
X[:, i, 1],
alpha=0.2,
marker='x',
label='Cluster {}'.format(i)
)
plot_covariance(
means[i, :],
covars[i, :, :],
ax,
color=sc.get_facecolor()[0]
)
sc = ax.scatter(
gmm.means[:, 0],
gmm.means[:, 1],
marker='+',
label='Fitted Gaussians'
)
for i in range(K):
plot_covariance(
gmm.means[i, :],
gmm.covars[i, :, :],
ax,
color=sc.get_facecolor()[0]
)
ax.legend()
plt.show()
weights, means, covars = torch.load(
'results/variable_k/em_128_508798_params.pkl',
map_location=torch.device('cpu'))
X = torch.distributions.MultivariateNormal(loc=means,
covariance_matrix=covars).sample(
(N, )).numpy()
j = torch.distributions.Categorical(probs=weights.squeeze()).sample(
(N, )).numpy()
X = X[np.arange(N), j, :]
X_train = [X[:N_train]]
X_val = [X[N_train:]]
# X_data = [torch.Tensor(X.reshape(-1, D).astype(np.float32))]
# X_val = [torch.Tensor(X_val.reshape(-1, D).astype(np.float32))]
#device = torch.device('cuda')
#torch.set_default_tensor_type('torch.cuda.FloatTensor')
# torch.multiprocessing.set_start_method('spawn')
m = AffineCouplingFlow(D, 10, 100)
m.fit(X_train, val_data=X_val)
gmm = SGDGMM(K, D, epochs=100, batch_size=256)
gmm.fit(X_train, val_data=X_val, verbose=True)
gmm_params.append(path)
elbo_params = collections.defaultdict(list)
for f in os.listdir(args.elbo_results_dir):
path = os.path.join(args.elbo_results_dir, f)
elbo_params[int(f[16:18])].append(path)
iw_params = collections.defaultdict(list)
for f in os.listdir(args.iw_results_dir):
path = os.path.join(args.iw_results_dir, f)
iw_params[int(f[16:18])].append(path)
gmm = SGDDeconvGMM(K, D, batch_size=200, device=torch.device('cuda'))
test_gmm = SGDGMM(K, D, batch_size=200, device=torch.device('cuda'))
svi = SVIFlow(2,
5,
device=torch.device('cuda'),
batch_size=512,
epochs=100,
lr=1e-4,
n_samples=50,
use_iwae=False,
context_size=64,
hidden_features=128)
results = []
for p in gmm_params:
import matplotlib
import matplotlib.pyplot as plt
import seaborn as sns
sns.set()
import corner
from deconv.flow.maf import MAFlow
from deconv.gmm.sgd_gmm import SGDGMM
torch.set_default_tensor_type(torch.FloatTensor)
N = 10000
N_train = 8000
gmm = SGDGMM(2, 2, batch_size=512, device=torch.device('cuda'))
gmm.module.means.data = torch.Tensor([[-10.0, 0.0],
[10.0, 0.0]]).to(torch.device('cuda'))
gmm.module.soft_weights.data = torch.Tensor([0.0,
0.0]).to(torch.device('cuda'))
torch.set_default_tensor_type(torch.cuda.FloatTensor)
flow_kl = MAFlow(2,
5,
lr=1e-3,
epochs=500,
batch_size=512,
device=torch.device('cuda'))