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')
def train_and_eval(flow, epochs, lr, train_loader, test_loader,
target_distribution):
print('no of parameters is',
sum(param.numel() for param in flow.parameters()))
optimizer = torch.optim.Adam(flow.parameters(), lr=lr)
train_losses, test_losses = [], []
for epoch in range(epochs):
print('Starting epoch:', epoch + 1, 'of', 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
if __name__ == '__main__':
print('Device is:', device)
from torch.distributions.normal import Normal
import numpy as np
flow = RealNVP(INPUT_H, INPUT_W).to(device)
target_distribution = Normal(
torch.tensor(0).float().to(device),
torch.tensor(1).float().to(device))
flow, train_losses, test_losses = train_and_eval(flow, 100, 5e-4,
train_loader, test_loader,
target_distribution)
print('train losses are', train_losses)
print('test losses are', test_losses)
torch.save(flow.state_dict(), 'trained_weights.pt')
