From 72e4bc5d20e153800f19a94e4bfd075adf30e3f3 Mon Sep 17 00:00:00 2001 From: Francois Fleuret Date: Thu, 19 Dec 2019 02:11:29 +0100 Subject: [PATCH] Cleaning up. Now saving all the figures. --- denoising-ae-field.py | 127 +++++++++++++++++++++--------------------- 1 file changed, 64 insertions(+), 63 deletions(-) diff --git a/denoising-ae-field.py b/denoising-ae-field.py index effee19..8f748d1 100755 --- a/denoising-ae-field.py +++ b/denoising-ae-field.py @@ -1,17 +1,10 @@ #!/usr/bin/env python import math +import matplotlib.pyplot as plt -import torch, torchvision - +import torch from torch import nn -from torch.nn import functional as F - -model = nn.Sequential( - nn.Linear(2, 100), - nn.ReLU(), - nn.Linear(100, 2) -) ###################################################################### @@ -22,14 +15,14 @@ def data_zigzag(nb): y = a * 2.5 - 1.25 data = torch.cat((y, x), 1) data = data @ torch.tensor([[1., -1.], [1., 1.]]) - return data + return data, 'zigzag' def data_spiral(nb): a = torch.empty(nb).uniform_(0, 1).view(-1, 1) x = (a * 2.25 * math.pi).cos() * (a * 0.8 + 0.5) y = (a * 2.25 * math.pi).sin() * (a * 0.8 + 0.5) data = torch.cat((y, x), 1) - return data + return data, 'spiral' def data_penta(nb): a = (torch.randint(5, (nb,)).float() / 5 * 2 * math.pi).view(-1, 1) @@ -37,70 +30,78 @@ def data_penta(nb): y = a.sin() data = torch.cat((y, x), 1) data = data + data.new(data.size()).normal_(0, 0.05) - return data + return data, 'penta' ###################################################################### -data = data_spiral(1000) -# data = data_zigzag(1000) -# data = data_penta(1000) - -data = data - data.mean(0) - -batch_size, nb_epochs = 100, 1000 -optimizer = torch.optim.Adam(model.parameters(), lr = 1e-3) -criterion = nn.MSELoss() - -for e in range(nb_epochs): - acc_loss = 0 - for input in data.split(batch_size): - noise = input.new(input.size()).normal_(0, 0.1) - output = model(input + noise) - loss = criterion(output, input) - acc_loss += loss.item() - optimizer.zero_grad() - loss.backward() - optimizer.step() - if (e+1)%10 == 0: print(e+1, acc_loss) +def train_model(data): + model = nn.Sequential( + nn.Linear(2, 100), + nn.ReLU(), + nn.Linear(100, 2) + ) + + batch_size, nb_epochs = 100, 1000 + optimizer = torch.optim.Adam(model.parameters(), lr = 1e-3) + criterion = nn.MSELoss() + + for e in range(nb_epochs): + acc_loss = 0 + for input in data.split(batch_size): + noise = input.new(input.size()).normal_(0, 0.1) + output = model(input + noise) + loss = criterion(output, input) + acc_loss += loss.item() + optimizer.zero_grad() + loss.backward() + optimizer.step() + if (e+1)%100 == 0: print(e+1, acc_loss) + + return model ###################################################################### -a = torch.linspace(-1.5, 1.5, 30) -x = a.view( 1, -1, 1).expand(a.size(0), a.size(0), 1) -y = a.view(-1, 1, 1).expand(a.size(0), a.size(0), 1) -grid = torch.cat((y, x), 2).view(-1, 2) +def save_image(data, data_name, model): + a = torch.linspace(-1.5, 1.5, 30) + x = a.view( 1, -1, 1).expand(a.size(0), a.size(0), 1) + y = a.view(-1, 1, 1).expand(a.size(0), a.size(0), 1) + grid = torch.cat((y, x), 2).view(-1, 2) -# Take the origins of the arrows on the part of grid closer than 0.1 -# from the data points -dist = (grid.view(-1, 1, 2) - data.view(1, -1, 2)).pow(2).sum(2).min(1)[0] -origins = grid[torch.arange(grid.size(0)).masked_select(dist < 0.1)] + # Take the origins of the arrows on the part of the grid closer than + # sqrt(0.1) to the data points + dist = (grid.view(-1, 1, 2) - data.view(1, -1, 2)).pow(2).sum(2).min(1)[0] + origins = grid[torch.arange(grid.size(0)).masked_select(dist < 0.1)] -field = model(origins).detach() - origins + field = model(origins).detach() - origins -###################################################################### - -import matplotlib.pyplot as plt + fig = plt.figure() + ax = fig.add_subplot(1, 1, 1) -fig = plt.figure() -ax = fig.add_subplot(1, 1, 1) + ax.axis('off') + ax.set_xlim(-1.6, 1.6) + ax.set_ylim(-1.6, 1.6) + ax.set_aspect(1) -ax.axis('off') -ax.set_xlim(-1.6, 1.6) -ax.set_ylim(-1.6, 1.6) -ax.set_aspect(1) + plot_field = ax.quiver( + origins[:, 0].numpy(), origins[:, 1].numpy(), + field[:, 0].numpy(), field[:, 1].numpy(), + units = 'xy', scale = 1, + width = 3e-3, headwidth = 25, headlength = 25 + ) -plot_field = ax.quiver( - origins[:, 0].numpy(), origins[:, 1].numpy(), - field[:, 0].numpy(), field[:, 1].numpy(), - units = 'xy', scale = 1, - width = 3e-3, headwidth = 25, headlength = 25 -) + plot_data = ax.scatter( + data[:, 0].numpy(), data[:, 1].numpy(), + s = 1, color = 'tab:blue' + ) -plot_data = ax.scatter( - data[:, 0].numpy(), data[:, 1].numpy(), - s = 1, color = 'tab:blue' -) - -fig.savefig('denoising_field.pdf', bbox_inches='tight') + filename = f'denoising_field_{data_name}.pdf' + print(f'Saving {filename}') + fig.savefig(filename, bbox_inches='tight') ###################################################################### + +for data_source in [ data_zigzag, data_spiral, data_penta ]: + data, data_name = data_source(1000) + data = data - data.mean(0) + model = train_model(data) + save_image(data, data_name, model) -- 2.20.1