def main(args):
scale_network = lambda: nn.Sequential(
nn.Linear(args.hidden_dim, args.hidden_size), nn.LeakyReLU(),
nn.Linear(args.hidden_size, args.hidden_size), nn.LeakyReLU(),
nn.Linear(args.hidden_size, args.hidden_dim), nn.Tanh())
translation_network = lambda: nn.Sequential(
nn.Linear(args.hidden_dim, args.hidden_size), nn.LeakyReLU(),
nn.Linear(args.hidden_size, args.hidden_size), nn.LeakyReLU(),
nn.Linear(args.hidden_size, args.hidden_dim))
masks = torch.tensor([1, 0] * args.num_layers, dtype=torch.float)
masks = torch.stack((masks, 1 - masks), dim=1)
prior = torch.distributions.MultivariateNormal(
torch.zeros(args.hidden_dim), torch.eye(args.hidden_dim))
flow = RealNVP(scale_network, translation_network, masks, prior)
optimizer = optim.Adam(
filter(lambda x: x.requires_grad == True, flow.parameters()))
for epoch in range(args.num_epoch):
loss = run_epoch(flow, optimizer, args.batch_size)
if epoch % args.log_train_step == 0:
logging.info(" Epoch: {} | Loss: {}".format(epoch, loss))
test(flow)
def train_and_eval(epochs, lr, train_loader, test_loader, target_distribution):
transforms = [
AffineTransform2D(True),
AffineTransform2D(False),
AffineTransform2D(True),
AffineTransform2D(False)
]
flow = RealNVP(transforms)
optimizer = torch.optim.Adam(flow.parameters(), lr=lr)
train_losses, test_losses = [], []
for epoch in range(epochs):
train(flow, train_loader, optimizer, target_distribution)
train_losses.append(eval_loss(flow, train_loader, target_distribution))
test_losses.append(eval_loss(flow, test_loader, target_distribution))
return flow, train_losses, test_losses
def train(param, x, y):
dim_in = x.shape[1]
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(device)
dataloader = DataLoader(torch.from_numpy(x.astype(np.float32)),
batch_size=param.batch_size,
shuffle=True,
num_workers=2)
flow = RealNVP(dim_in, device)
flow.to(device)
flow.train()
optimizer = torch.optim.Adam(
[p for p in flow.parameters() if p.requires_grad == True], lr=param.lr)
it, print_cnt = 0, 0
while it < param.total_it:
for i, data in enumerate(dataloader):
loss = -flow.log_prob(data.to(device)).mean()
optimizer.zero_grad()
loss.backward(retain_graph=True)
optimizer.step()
it += data.shape[0]
print_cnt += data.shape[0]
if print_cnt > PRINT_FREQ:
print('it {:d} -- loss {:.03f}'.format(it, loss))
print_cnt = 0
torch.save(flow.state_dict(), 'flow_model.pytorch')
train_loader = torch.utils.data.DataLoader(train_data,
batch_size=batch_size,
shuffle=True,
**kwargs)
# Set figures
plt.subplots(nrows=2, ncols=2)
plt.subplots_adjust(hspace=0.5, wspace=0.3)
plt.subplot(2, 2, 1)
plt.scatter(input[:, 0], input[:, 1], c='b', s=10)
plt.title("INPUT: x ~ p(x)")
# Set model
mask = torch.tensor([0.0, 1.0])
model = RealNVP(INPUT_DIM, OUTPUT_DIM, HID_DIM, mask, 8)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
log_gaussian = torch.distributions.MultivariateNormal(torch.zeros(2),
torch.eye(2))
def train(args):
"Forward flow data for construction normalization distribution"
model.train()
for epoch in range(args.epochs):
train_loss = 0.0
for i, data in enumerate(train_loader):
optimizer.zero_grad()
z, log_det_sum = model(data)
loss = -(log_gaussian.log_prob(z.float()) + log_det_sum).mean()
