def test_input_extra_var(self):
q = Normal(var=['z'], cond_var=['x'], loc='x', scale=1)
p = Normal(var=['y'], cond_var=['z'], loc='z', scale=1)
e = Expectation(q, p.log_prob())
assert set(e.eval({'y': torch.zeros(1), 'x': torch.zeros(1),
'w': torch.zeros(1)}, return_dict=True)[1]) == set(('w', 'x', 'y', 'z'))
assert set(e.eval({'y': torch.zeros(1), 'x': torch.zeros(1),
'z': torch.zeros(1)}, return_dict=True)[1]) == set(('x', 'y', 'z'))
def test_input_extra_var(self):
normal = Normal(loc=0, scale=1)
assert set(normal.sample({'y': torch.zeros(1)})) == set(('x', 'y'))
assert normal.get_log_prob({
'y': torch.zeros(1),
'x': torch.zeros(1)
}).shape == torch.Size([1])
assert set(normal.sample({'x': torch.zeros(1)})) == set(('x'))
def test_get_entropy(self):
dist = Normal(var=['x'], cond_var=['y'], loc='y', scale=1)
truth = dist.get_entropy({'y': torch.ones(1)})
dist = Normal(var=['x'], cond_var=['y'], loc='y',
scale=1).replace_var(y='z', x='y')
result = dist.get_entropy({'z': torch.ones(1)})
assert result == truth
dist = Normal(var=['x'], cond_var=['y'], loc='y',
scale=1).replace_var(y='z')
with pytest.raises(ValueError):
dist.get_entropy({'y': torch.ones(1)})
def test_sample_variance(self):
dist = Normal(var=['x'], cond_var=['y'], loc=2, scale='y')
result = dist.sample_variance({'y': torch.ones(1)})
assert result == torch.ones(1)
dist = Normal(var=['x'], cond_var=['y'], loc=2,
scale='y').replace_var(y='z')
result = dist.sample_variance({'z': torch.ones(1)})
assert result == torch.ones(1)
dist = Normal(var=['x'], cond_var=['y'], loc=2,
scale='y').replace_var(y='z')
with pytest.raises(ValueError):
dist.sample_variance({'y': torch.ones(1)})
def test_save_dist(tmpdir, no_contiguous_tensor):
# pull request:#110
ones = torch.ones_like(no_contiguous_tensor)
p = Normal(loc=no_contiguous_tensor, scale=ones)
save_path = pjoin(tmpdir, "tmp.pt")
torch.save(p.state_dict(), save_path)
q = Normal(loc=ones, scale=3 * ones)
assert not torch.all(no_contiguous_tensor == q.loc).item()
# it needs copy of tensor
q = Normal(loc=ones, scale=ones)
q.load_state_dict(torch.load(save_path))
assert torch.all(no_contiguous_tensor == q.loc).item()
def test_init_with_scalar_params(self):
normal = Normal(loc=0, scale=1, features_shape=[2])
assert normal.sample()['x'].shape == torch.Size([1, 2])
assert normal.features_shape == torch.Size([2])
normal = Normal(loc=0, scale=1)
assert normal.sample()['x'].shape == torch.Size([1])
assert normal.features_shape == torch.Size([])
def test_set_option(self):
dist = Normal(var=['x'], cond_var=['y'], loc='y', scale=1) * Normal(
var=['y'], loc=0, scale=1)
dist.graph.set_option(dict(batch_n=4, sample_shape=(2, 3)), ['y'])
sample = dist.sample()
assert sample['y'].shape == torch.Size([2, 3, 4])
assert sample['x'].shape == torch.Size([2, 3, 4])
dist.graph.set_option({}, ['y'])
assert dist.get_log_prob(sample, sum_features=True,
feature_dims=None).shape == torch.Size([2])
assert dist.get_log_prob(
sample, sum_features=False).shape == torch.Size([2, 3, 4])
dist = Normal(var=['x'], cond_var=['y'], loc='y',
scale=1) * FactorizedBernoulli(
var=['y'], probs=torch.tensor([0.3, 0.8]))
dist.graph.set_option(dict(batch_n=3, sample_shape=(4, )), ['y'])
sample = dist.sample()
assert sample['y'].shape == torch.Size([4, 3, 2])
assert sample['x'].shape == torch.Size([4, 3, 2])
dist.graph.set_option(dict(), ['y'])
assert dist.get_log_prob(sample, sum_features=True,
feature_dims=[-1]).shape == torch.Size([4, 3])
def test_get_log_prob_feature_dims2(self):
dist = Normal(var=['x'], cond_var=['y'], loc='y', scale=1) * Normal(
var=['y'], loc=0, scale=1)
dist.graph.set_option(dict(batch_n=4, sample_shape=(2, 3)), ['y'])
sample = dist.sample()
assert sample['y'].shape == torch.Size([2, 3, 4])
list(dist.graph._factors_from_variable('y'))[0].option = {}
assert dist.get_log_prob(sample, sum_features=True,
feature_dims=None).shape == torch.Size([2])
assert dist.get_log_prob(sample, sum_features=True,
feature_dims=[-2]).shape == torch.Size([2, 4])
assert dist.get_log_prob(sample,
sum_features=True,
feature_dims=[0, 1]).shape == torch.Size([4])
assert dist.get_log_prob(sample, sum_features=True,
feature_dims=[]).shape == torch.Size(
[2, 3, 4])
def test_sample_mean(self):
dist = MixtureModel([Normal(loc=0, scale=1),
Normal(loc=1, scale=1)],
Categorical(probs=torch.tensor([1., 2.])))
assert dist.sample(sample_mean=True)['x'] == torch.ones(1)
def test_batch_n(self):
normal = Normal(loc=0, scale=1)
assert normal.sample(batch_n=3)['x'].shape == torch.Size([3])
def main(smoke_test=False):
epochs = 2 if smoke_test == True else 50
batch_size = 128
seed = 0
x_ch = 1
z_dim = 32
# 乱数シード初期化
torch.manual_seed(seed)
torch.random.manual_seed(seed)
torch.cuda.manual_seed(seed)
date_and_time = datetime.datetime.now().strftime('%Y-%m%d-%H%M')
save_root = f'./results/pixyz/{date_and_time}'
if not os.path.exists(save_root):
os.makedirs(save_root)
if torch.cuda.is_available():
device = 'cuda:0'
else:
device = 'cpu'
root = '/mnt/hdd/sika/Datasets'
train_loader, test_loader = make_MNIST_loader(root, batch_size=batch_size)
# 生成モデルと推論モデルの生成
p = Generator(x_ch, z_dim).to(device)
q = Inference(x_ch, z_dim).to(device)
# 潜在変数の事前分布の規定
p_prior = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.),
var=['z'], features_shape=[z_dim], name='p_{prior}').to(device)
# 損失関数の定義
loss = (KullbackLeibler(q, p_prior) - Expectation(q, LogProb(p))).mean()
# Model APIの設定
model = Model(loss=loss, distributions=[p, q],
optimizer=optim.Adam, optimizer_params={"lr": 1e-3})
x_fixed, _ = next(iter(test_loader))
x_fixed = x_fixed[:8].to(device)
z_fixed = p_prior.sample(batch_n=64)['z']
train_loss_list, test_loss_list = [], []
for epoch in range(1, epochs + 1):
train_loss_list.append(learn(model, epoch, train_loader, device, train=True))
test_loss_list.append(learn(model, epoch, test_loader, device, train=False))
print(f' [Epoch {epoch}] train loss {train_loss_list[-1]:.4f}')
print(f' [Epoch {epoch}] test loss {test_loss_list[-1]:.4f}\n')
# 損失値のグラフを作成し保存
plt.plot(list(range(1, epoch+1)), train_loss_list, label='train')
plt.plot(list(range(1, epoch + 1)), test_loss_list, label='test')
plt.xlabel('epochs')
plt.ylabel('loss')
plt.legend()
plt.savefig(os.path.join(save_root, 'loss.png'))
plt.close()
# 再構成画像
x_reconst = reconstruct_image(p, q, x_fixed)
save_image(torch.cat([x_fixed, x_reconst], dim=0), os.path.join(save_root, f'reconst_{epoch}.png'), nrow=8)
# 補間画像
x_interpol = interpolate_image(p, q, x_fixed)
save_image(x_interpol, os.path.join(save_root, f'interpol_{epoch}.png'), nrow=8)
# 生成画像（潜在変数固定）
x_generate = generate_image(p, z_fixed)
save_image(x_generate, os.path.join(save_root, f'generate_{epoch}.png'), nrow=8)
# 生成画像（ランダムサンプリング）
x_sample = sample_image(p_prior, p)
save_image(x_sample, os.path.join(save_root, f'sample_{epoch}.png'), nrow=8)
class TestDistributionBase:
def test_init_with_scalar_params(self):
normal = Normal(loc=0, scale=1, features_shape=[2])
assert normal.sample()['x'].shape == torch.Size([1, 2])
assert normal.features_shape == torch.Size([2])
normal = Normal(loc=0, scale=1)
assert normal.sample()['x'].shape == torch.Size([1])
assert normal.features_shape == torch.Size([])
def test_batch_n(self):
normal = Normal(loc=0, scale=1)
assert normal.sample(batch_n=3)['x'].shape == torch.Size([3])
def test_input_extra_var(self):
normal = Normal(loc=0, scale=1)
assert set(normal.sample({'y': torch.zeros(1)})) == set(('x', 'y'))
assert normal.get_log_prob({
'y': torch.zeros(1),
'x': torch.zeros(1)
}).shape == torch.Size([1])
assert set(normal.sample({'x': torch.zeros(1)})) == set(('x'))
def test_sample_mean(self):
dist = Normal(loc=0, scale=1)
assert dist.sample(sample_mean=True)['x'] == torch.zeros(1)
@pytest.mark.parametrize(
"dist",
[
Normal(loc=0, scale=1),
Normal(var=['x'], loc=0, scale=1) *
Normal(var=['y'], loc=0, scale=1),
# Normal(var=['x'], cond_var=['y'], loc='y', scale=1) * Normal(var=['y'], loc=0, scale=1),
],
)
def test_get_log_prob_feature_dims(self, dist):
assert dist.get_log_prob(dist.sample(batch_n=4, sample_shape=(2, 3)),
sum_features=True,
feature_dims=None).shape == torch.Size([2])
assert dist.get_log_prob(dist.sample(batch_n=4, sample_shape=(2, 3)),
sum_features=True,
feature_dims=[-2]).shape == torch.Size([2, 4])
assert dist.get_log_prob(dist.sample(batch_n=4, sample_shape=(2, 3)),
sum_features=True,
feature_dims=[0, 1]).shape == torch.Size([4])
assert dist.get_log_prob(dist.sample(batch_n=4, sample_shape=(2, 3)),
sum_features=True,
feature_dims=[]).shape == torch.Size(
[2, 3, 4])
def test_get_log_prob_feature_dims2(self):
dist = Normal(var=['x'], cond_var=['y'], loc='y', scale=1) * Normal(
var=['y'], loc=0, scale=1)
dist.graph.set_option(dict(batch_n=4, sample_shape=(2, 3)), ['y'])
sample = dist.sample()
assert sample['y'].shape == torch.Size([2, 3, 4])
list(dist.graph._factors_from_variable('y'))[0].option = {}
assert dist.get_log_prob(sample, sum_features=True,
feature_dims=None).shape == torch.Size([2])
assert dist.get_log_prob(sample, sum_features=True,
feature_dims=[-2]).shape == torch.Size([2, 4])
assert dist.get_log_prob(sample,
sum_features=True,
feature_dims=[0, 1]).shape == torch.Size([4])
assert dist.get_log_prob(sample, sum_features=True,
feature_dims=[]).shape == torch.Size(
[2, 3, 4])
@pytest.mark.parametrize("dist", [
Normal(loc=0, scale=1),
Normal(var=['x'], cond_var=['y'], loc='y', scale=1) *
Normal(var=['y'], loc=0, scale=1),
])
def test_unknown_option(self, dist):
x_dict = dist.sample(unknown_opt=None)
dist.get_log_prob(x_dict, unknown_opt=None)
def test_sample_mean(self):
dist = Normal(var=['x'], loc=0, scale=1) * Normal(
var=['y'], cond_var=['x'], loc='x', scale=1)
assert dist.sample(sample_mean=True)['y'] == torch.zeros(1)
def test_rename_atomdist(self):
normal = Normal(var=['x'], name='p')
graph = normal.graph
assert graph.name == 'p'
normal.name = 'q'
assert graph.name == 'q'
def test_sample_mean(self):
p = Normal(loc=0, scale=1)
f = p.log_prob()
e = Expectation(p, f)
e.eval({}, sample_mean=True)
def test_sample_mean(self):
dist = Normal(loc=0, scale=1)
assert dist.sample(sample_mean=True)['x'] == torch.zeros(1)
def test_get_params(self):
dist = Normal(var=['x'], cond_var=['y'], loc='y', scale=1)
result = dist.get_params({'y': torch.ones(1)})
assert list(result.keys()) == ['loc', 'scale']
dist = Normal(var=['x'], cond_var=['y'], loc='y',
scale=1).replace_var(y='z')
result = dist.get_params({'z': torch.ones(1)})
assert list(result.keys()) == ['loc', 'scale']
dist = Normal(var=['x'], cond_var=['y'], loc='y',
scale=1).replace_var(y='z')
with pytest.raises(ValueError):
dist.get_params({'y': torch.ones(1)})
dist = Normal(var=['x'], cond_var=['y'], loc='y',
scale=1).replace_var(x='z')
result = dist.get_params({'y': torch.ones(1)})
assert list(result.keys()) == ['loc', 'scale']
dist = Normal(var=['x'], cond_var=['y'], loc='y', scale=1) * Normal(
var=['y'], loc=0, scale=1)
with pytest.raises(NotImplementedError):
dist.get_params()
def test_print(self):
normal = Normal(var=['x'], name='p')
print(normal.graph)
def test_input_var(self):
q = Normal(var=['z'], cond_var=['x'], loc='x', scale=1)
p = Normal(var=['y'], cond_var=['z'], loc='z', scale=1)
e = Expectation(q, p.log_prob())
assert set(e.input_var) == set(('x', 'y'))
assert e.eval({'y': torch.zeros(1), 'x': torch.zeros(1)}).shape == torch.Size([1])
def prior():
return Normal(var=["z"],
name="p_{prior}",
features_shape=[10],
loc=torch.tensor(0.),
scale=torch.tensor(1.))
class Classifier(RelaxedCategorical):
def __init__(self):
super(Classifier, self).__init__(var=["y"], cond_var=["x"], name="p")
self.fc1 = nn.Linear(x_dim, 512)
self.fc2 = nn.Linear(512, y_dim)
def forward(self, x):
h = F.relu(self.fc1(x))
h = F.softmax(self.fc2(h), dim=1)
return {"probs": h}
# prior model p(z)
prior = Normal(loc=torch.tensor(0.),
scale=torch.tensor(1.),
var=["z"],
features_shape=[z_dim],
name="p_{prior}").to(device)
# distributions for supervised learning
p = Generator().to(device)
q = Inference().to(device)
f = Classifier().to(device)
p_joint = p * prior
# distributions for unsupervised learning
_q_u = q.replace_var(x="x_u", y="y_u")
p_u = p.replace_var(x="x_u", y="y_u")
f_u = f.replace_var(x="x_u", y="y_u")
q_u = _q_u * f_u
if torch.cuda.is_available():
device = "cuda"
else:
device = "cpu"
# p(y|x,z)
p = net.Generator_().to(device)
# q(z|x,y)
q = net.Inference().to(device)
# prior p(z)
prior = Normal(loc=torch.tensor(0.0),
scale=torch.tensor(1.0),
var=["z"],
features_shape=[net.Z_DIM],
name="p_{prior}").to(device)
loss = (KullbackLeibler(q, prior) - Expectation(q, LogProb(p))).mean()
model = Model(loss=loss,
distributions=[p, q],
optimizer=optim.Adam,
optimizer_params={"lr": 1e-3})
# print(model)
x_org, y_org = next(iter(test_loader))
# 再構築用サンプルデータ
x_fixed = x_org[:8].to(device)
y_fixed = y_org[:8]