From 05b9b133a45ac9bd5abe6f8b6d29095f9c82797a Mon Sep 17 00:00:00 2001 From: =?utf8?q?Fran=C3=A7ois=20Fleuret?= Date: Sun, 14 May 2023 22:22:34 +0200 Subject: [PATCH] Update. --- ae_size.py | 7 +- attentiontoy1d.py | 292 ++++++++++++++++++++++++--------------- causal-autoregression.py | 214 ++++++++++++++-------------- confidence.py | 28 ++-- conv_chain.py | 25 ++-- ddpol.py | 141 +++++++++++-------- denoising-ae-field.py | 67 +++++---- elbo.py | 14 +- flatparam.py | 25 ++-- gpt-test.py | 37 +++-- hallu.py | 20 +-- lazy_linear.py | 24 ++-- mandelbrot.py | 2 +- mi_estimator.py | 239 ++++++++++++++++++-------------- minidiffusion.py | 220 +++++++++++++++-------------- miniflow.py | 84 ++++++----- poly.py | 12 +- rmax.py | 40 ++++++ sizer.py | 21 +-- speed.py | 22 +-- tensorstack.py | 51 +++---- tinyae.py | 91 ++++++------ tinymnist.py | 36 ++--- 23 files changed, 987 insertions(+), 725 deletions(-) create mode 100755 rmax.py diff --git a/ae_size.py b/ae_size.py index 067a7fa..49f4a20 100755 --- a/ae_size.py +++ b/ae_size.py @@ -11,8 +11,9 @@ from torch import Tensor ###################################################################### + def minimal_input_size(w, layer_specs): - assert w > 0, 'The input is too small' + assert w > 0, "The input is too small" if layer_specs == []: return w else: @@ -21,13 +22,13 @@ def minimal_input_size(w, layer_specs): v = minimal_input_size(v, layer_specs[1:]) return (v - 1) * stride + kernel_size + ###################################################################### # Dummy test if __name__ == "__main__": - - layer_specs = [ (17, 5), (5, 4), (3, 2), (3, 2) ] + layer_specs = [(17, 5), (5, 4), (3, 2), (3, 2)] layers = [] diff --git a/attentiontoy1d.py b/attentiontoy1d.py index b463340..d2db9c6 100755 --- a/attentiontoy1d.py +++ b/attentiontoy1d.py @@ -14,26 +14,34 @@ import matplotlib.pyplot as plt ###################################################################### -parser = argparse.ArgumentParser(description='Toy attention model.') - -parser.add_argument('--nb_epochs', - type = int, default = 250) - -parser.add_argument('--with_attention', - help = 'Use the model with an attention layer', - action='store_true', default=False) - -parser.add_argument('--group_by_locations', - help = 'Use the task where the grouping is location-based', - action='store_true', default=False) - -parser.add_argument('--positional_encoding', - help = 'Provide a positional encoding', - action='store_true', default=False) - -parser.add_argument('--seed', - type = int, default = 0, - help = 'Random seed (default 0, < 0 is no seeding)') +parser = argparse.ArgumentParser(description="Toy attention model.") + +parser.add_argument("--nb_epochs", type=int, default=250) + +parser.add_argument( + "--with_attention", + help="Use the model with an attention layer", + action="store_true", + default=False, +) + +parser.add_argument( + "--group_by_locations", + help="Use the task where the grouping is location-based", + action="store_true", + default=False, +) + +parser.add_argument( + "--positional_encoding", + help="Provide a positional encoding", + action="store_true", + default=False, +) + +parser.add_argument( + "--seed", type=int, default=0, help="Random seed (default 0, < 0 is no seeding)" +) args = parser.parse_args() @@ -42,32 +50,37 @@ if args.seed >= 0: ###################################################################### -label='' +label = "" -if args.with_attention: label = 'wa_' +if args.with_attention: + label = "wa_" -if args.group_by_locations: label += 'lg_' +if args.group_by_locations: + label += "lg_" -if args.positional_encoding: label += 'pe_' +if args.positional_encoding: + label += "pe_" -log_file = open(f'att1d_{label}train.log', 'w') +log_file = open(f"att1d_{label}train.log", "w") ###################################################################### + def log_string(s): if log_file is not None: - log_file.write(s + '\n') + log_file.write(s + "\n") log_file.flush() print(s) sys.stdout.flush() + ###################################################################### if torch.cuda.is_available(): - device = torch.device('cuda') + device = torch.device("cuda") torch.backends.cudnn.benchmark = True else: - device = torch.device('cpu') + device = torch.device("cpu") ###################################################################### @@ -75,37 +88,51 @@ seq_height_min, seq_height_max = 1.0, 25.0 seq_width_min, seq_width_max = 5.0, 11.0 seq_length = 100 -def positions_to_sequences(tr = None, bx = None, noise_level = 0.3): - st = torch.arange(seq_length, device = device).float() + +def positions_to_sequences(tr=None, bx=None, noise_level=0.3): + st = torch.arange(seq_length, device=device).float() st = st[None, :, None] tr = tr[:, None, :, :] bx = bx[:, None, :, :] - xtr = torch.relu(tr[..., 1] - torch.relu(torch.abs(st - tr[..., 0]) - 0.5) * 2 * tr[..., 1] / tr[..., 2]) - xbx = torch.sign(torch.relu(bx[..., 1] - torch.abs((st - bx[..., 0]) * 2 * bx[..., 1] / bx[..., 2]))) * bx[..., 1] + xtr = torch.relu( + tr[..., 1] + - torch.relu(torch.abs(st - tr[..., 0]) - 0.5) * 2 * tr[..., 1] / tr[..., 2] + ) + xbx = ( + torch.sign( + torch.relu( + bx[..., 1] - torch.abs((st - bx[..., 0]) * 2 * bx[..., 1] / bx[..., 2]) + ) + ) + * bx[..., 1] + ) x = torch.cat((xtr, xbx), 2) - u = F.max_pool1d(x.sign().permute(0, 2, 1), kernel_size = 2, stride = 1).permute(0, 2, 1) + u = F.max_pool1d(x.sign().permute(0, 2, 1), kernel_size=2, stride=1).permute( + 0, 2, 1 + ) collisions = (u.sum(2) > 1).max(1).values y = x.max(2).values return y + torch.rand_like(y) * noise_level - noise_level / 2, collisions + ###################################################################### -def generate_sequences(nb): +def generate_sequences(nb): # Position / height / width - tr = torch.empty(nb, 2, 3, device = device) - tr[:, :, 0].uniform_(seq_width_max/2, seq_length - seq_width_max/2) + tr = torch.empty(nb, 2, 3, device=device) + tr[:, :, 0].uniform_(seq_width_max / 2, seq_length - seq_width_max / 2) tr[:, :, 1].uniform_(seq_height_min, seq_height_max) tr[:, :, 2].uniform_(seq_width_min, seq_width_max) - bx = torch.empty(nb, 2, 3, device = device) - bx[:, :, 0].uniform_(seq_width_max/2, seq_length - seq_width_max/2) + bx = torch.empty(nb, 2, 3, device=device) + bx[:, :, 0].uniform_(seq_width_max / 2, seq_length - seq_width_max / 2) bx[:, :, 1].uniform_(seq_height_min, seq_height_max) bx[:, :, 2].uniform_(seq_width_min, seq_width_max) @@ -117,7 +144,9 @@ def generate_sequences(nb): h_right = (a[:, :, 1] * (1 - mask_left)).sum(1) / 2 valid = (h_left - h_right).abs() > 4 else: - valid = (torch.abs(tr[:, 0, 1] - tr[:, 1, 1]) > 4) & (torch.abs(tr[:, 0, 1] - tr[:, 1, 1]) > 4) + valid = (torch.abs(tr[:, 0, 1] - tr[:, 1, 1]) > 4) & ( + torch.abs(tr[:, 0, 1] - tr[:, 1, 1]) > 4 + ) input, collisions = positions_to_sequences(tr, bx) @@ -125,13 +154,13 @@ def generate_sequences(nb): a = torch.cat((tr, bx), 1) v = a[:, :, 0].sort(1).values[:, 2:3] mask_left = (a[:, :, 0] < v).float() - h_left = (a[:, :, 1] * mask_left).sum(1, keepdim = True) / 2 - h_right = (a[:, :, 1] * (1 - mask_left)).sum(1, keepdim = True) / 2 + h_left = (a[:, :, 1] * mask_left).sum(1, keepdim=True) / 2 + h_right = (a[:, :, 1] * (1 - mask_left)).sum(1, keepdim=True) / 2 a[:, :, 1] = mask_left * h_left + (1 - mask_left) * h_right tr, bx = a.split(2, 1) else: - tr[:, :, 1:2] = tr[:, :, 1:2].mean(1, keepdim = True) - bx[:, :, 1:2] = bx[:, :, 1:2].mean(1, keepdim = True) + tr[:, :, 1:2] = tr[:, :, 1:2].mean(1, keepdim=True) + bx[:, :, 1:2] = bx[:, :, 1:2].mean(1, keepdim=True) targets, _ = positions_to_sequences(tr, bx) @@ -150,9 +179,11 @@ def generate_sequences(nb): return input, targets, tr, bx + ###################################################################### -def save_sequence_images(filename, sequences, tr = None, bx = None): + +def save_sequence_images(filename, sequences, tr=None, bx=None): fig = plt.figure() ax = fig.add_subplot(1, 1, 1) @@ -160,54 +191,69 @@ def save_sequence_images(filename, sequences, tr = None, bx = None): ax.set_ylim(-1, seq_height_max + 4) for u in sequences: - ax.plot( - torch.arange(u[0].size(0)) + 0.5, u[0], color = u[1], label = u[2] - ) + ax.plot(torch.arange(u[0].size(0)) + 0.5, u[0], color=u[1], label=u[2]) - ax.legend(frameon = False, loc = 'upper left') + ax.legend(frameon=False, loc="upper left") - delta = -1. + delta = -1.0 if tr is not None: - ax.scatter(tr[:, 0].cpu(), torch.full((tr.size(0),), delta), color = 'black', marker = '^', clip_on=False) + ax.scatter( + tr[:, 0].cpu(), + torch.full((tr.size(0),), delta), + color="black", + marker="^", + clip_on=False, + ) if bx is not None: - ax.scatter(bx[:, 0].cpu(), torch.full((bx.size(0),), delta), color = 'black', marker = 's', clip_on=False) + ax.scatter( + bx[:, 0].cpu(), + torch.full((bx.size(0),), delta), + color="black", + marker="s", + clip_on=False, + ) + + fig.savefig(filename, bbox_inches="tight") - fig.savefig(filename, bbox_inches='tight') + plt.close("all") - plt.close('all') ###################################################################### + class AttentionLayer(nn.Module): def __init__(self, in_channels, out_channels, key_channels): super().__init__() - self.conv_Q = nn.Conv1d(in_channels, key_channels, kernel_size = 1, bias = False) - self.conv_K = nn.Conv1d(in_channels, key_channels, kernel_size = 1, bias = False) - self.conv_V = nn.Conv1d(in_channels, out_channels, kernel_size = 1, bias = False) + self.conv_Q = nn.Conv1d(in_channels, key_channels, kernel_size=1, bias=False) + self.conv_K = nn.Conv1d(in_channels, key_channels, kernel_size=1, bias=False) + self.conv_V = nn.Conv1d(in_channels, out_channels, kernel_size=1, bias=False) def forward(self, x): Q = self.conv_Q(x) K = self.conv_K(x) V = self.conv_V(x) - A = einsum('nct,ncs->nts', Q, K).softmax(2) - y = einsum('nts,ncs->nct', A, V) + A = einsum("nct,ncs->nts", Q, K).softmax(2) + y = einsum("nts,ncs->nct", A, V) return y def __repr__(self): - return self._get_name() + \ - '(in_channels={}, out_channels={}, key_channels={})'.format( + return ( + self._get_name() + + "(in_channels={}, out_channels={}, key_channels={})".format( self.conv_Q.in_channels, self.conv_V.out_channels, - self.conv_K.out_channels + self.conv_K.out_channels, ) + ) def attention(self, x): Q = self.conv_Q(x) K = self.conv_K(x) - A = einsum('nct,ncs->nts', Q, K).softmax(2) + A = einsum("nct,ncs->nts", Q, K).softmax(2) return A + ###################################################################### train_input, train_targets, train_tr, train_bx = generate_sequences(25000) @@ -220,7 +266,9 @@ nc = 64 if args.positional_encoding: c = math.ceil(math.log(seq_length) / math.log(2.0)) - positional_input = (torch.arange(seq_length).unsqueeze(0) // 2**torch.arange(c).unsqueeze(1))%2 + positional_input = ( + torch.arange(seq_length).unsqueeze(0) // 2 ** torch.arange(c).unsqueeze(1) + ) % 2 positional_input = positional_input.unsqueeze(0).float() else: positional_input = torch.zeros(1, 0, seq_length) @@ -228,43 +276,41 @@ else: in_channels = 1 + positional_input.size(1) if args.with_attention: - model = nn.Sequential( - nn.Conv1d(in_channels, nc, kernel_size = ks, padding = ks//2), + nn.Conv1d(in_channels, nc, kernel_size=ks, padding=ks // 2), nn.ReLU(), - nn.Conv1d(nc, nc, kernel_size = ks, padding = ks//2), + nn.Conv1d(nc, nc, kernel_size=ks, padding=ks // 2), nn.ReLU(), AttentionLayer(nc, nc, nc), - nn.Conv1d(nc, nc, kernel_size = ks, padding = ks//2), + nn.Conv1d(nc, nc, kernel_size=ks, padding=ks // 2), nn.ReLU(), - nn.Conv1d(nc, 1, kernel_size = ks, padding = ks//2) + nn.Conv1d(nc, 1, kernel_size=ks, padding=ks // 2), ) else: - model = nn.Sequential( - nn.Conv1d(in_channels, nc, kernel_size = ks, padding = ks//2), + nn.Conv1d(in_channels, nc, kernel_size=ks, padding=ks // 2), nn.ReLU(), - nn.Conv1d(nc, nc, kernel_size = ks, padding = ks//2), + nn.Conv1d(nc, nc, kernel_size=ks, padding=ks // 2), nn.ReLU(), - nn.Conv1d(nc, nc, kernel_size = ks, padding = ks//2), + nn.Conv1d(nc, nc, kernel_size=ks, padding=ks // 2), nn.ReLU(), - nn.Conv1d(nc, nc, kernel_size = ks, padding = ks//2), + nn.Conv1d(nc, nc, kernel_size=ks, padding=ks // 2), nn.ReLU(), - nn.Conv1d(nc, 1, kernel_size = ks, padding = ks//2) + nn.Conv1d(nc, 1, kernel_size=ks, padding=ks // 2), ) nb_parameters = sum(p.numel() for p in model.parameters()) -with open(f'att1d_{label}model.log', 'w') as f: - f.write(str(model) + '\n\n') - f.write(f'nb_parameters {nb_parameters}\n') +with open(f"att1d_{label}model.log", "w") as f: + f.write(str(model) + "\n\n") + f.write(f"nb_parameters {nb_parameters}\n") ###################################################################### batch_size = 100 -optimizer = torch.optim.Adam(model.parameters(), lr = 1e-3) +optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) mse_loss = nn.MSELoss() model.to(device) @@ -278,9 +324,9 @@ mu, std = train_input.mean(), train_input.std() for e in range(args.nb_epochs): acc_loss = 0.0 - for input, targets in zip(train_input.split(batch_size), - train_targets.split(batch_size)): - + for input, targets in zip( + train_input.split(batch_size), train_targets.split(batch_size) + ): input = torch.cat((input, positional_input.expand(input.size(0), -1, -1)), 1) output = model((input - mu) / std) @@ -292,56 +338,58 @@ for e in range(args.nb_epochs): acc_loss += loss.item() - log_string(f'{e+1} {acc_loss}') + log_string(f"{e+1} {acc_loss}") ###################################################################### -train_input = train_input.detach().to('cpu') -train_targets = train_targets.detach().to('cpu') +train_input = train_input.detach().to("cpu") +train_targets = train_targets.detach().to("cpu") for k in range(15): save_sequence_images( - f'att1d_{label}train_{k:03d}.pdf', + f"att1d_{label}train_{k:03d}.pdf", [ - ( train_input[k, 0], 'blue', 'Input' ), - ( train_targets[k, 0], 'red', 'Target' ), + (train_input[k, 0], "blue", "Input"), + (train_targets[k, 0], "red", "Target"), ], ) #################### -test_input = torch.cat((test_input, positional_input.expand(test_input.size(0), -1, -1)), 1) +test_input = torch.cat( + (test_input, positional_input.expand(test_input.size(0), -1, -1)), 1 +) test_outputs = model((test_input - mu) / std).detach() if args.with_attention: k = next(k for k, l in enumerate(model) if isinstance(l, AttentionLayer)) x = model[0:k]((test_input - mu) / std) test_A = model[k].attention(x) - test_A = test_A.detach().to('cpu') + test_A = test_A.detach().to("cpu") -test_input = test_input.detach().to('cpu') -test_outputs = test_outputs.detach().to('cpu') -test_targets = test_targets.detach().to('cpu') -test_bx = test_bx.detach().to('cpu') -test_tr = test_tr.detach().to('cpu') +test_input = test_input.detach().to("cpu") +test_outputs = test_outputs.detach().to("cpu") +test_targets = test_targets.detach().to("cpu") +test_bx = test_bx.detach().to("cpu") +test_tr = test_tr.detach().to("cpu") for k in range(15): save_sequence_images( - f'att1d_{label}test_Y_{k:03d}.pdf', + f"att1d_{label}test_Y_{k:03d}.pdf", [ - ( test_input[k, 0], 'blue', 'Input' ), - ( test_outputs[k, 0], 'orange', 'Output' ), - ] + (test_input[k, 0], "blue", "Input"), + (test_outputs[k, 0], "orange", "Output"), + ], ) save_sequence_images( - f'att1d_{label}test_Yp_{k:03d}.pdf', + f"att1d_{label}test_Yp_{k:03d}.pdf", [ - ( test_input[k, 0], 'blue', 'Input' ), - ( test_outputs[k, 0], 'orange', 'Output' ), + (test_input[k, 0], "blue", "Input"), + (test_outputs[k, 0], "orange", "Output"), ], test_tr[k], - test_bx[k] + test_bx[k], ) if args.with_attention: @@ -350,15 +398,39 @@ for k in range(15): ax.set_xlim(0, seq_length) ax.set_ylim(0, seq_length) - ax.imshow(test_A[k], cmap = 'binary', interpolation='nearest') - delta = 0. - ax.scatter(test_bx[k, :, 0], torch.full((test_bx.size(1),), delta), color = 'black', marker = 's', clip_on=False) - ax.scatter(torch.full((test_bx.size(1),), delta), test_bx[k, :, 0], color = 'black', marker = 's', clip_on=False) - ax.scatter(test_tr[k, :, 0], torch.full((test_tr.size(1),), delta), color = 'black', marker = '^', clip_on=False) - ax.scatter(torch.full((test_tr.size(1),), delta), test_tr[k, :, 0], color = 'black', marker = '^', clip_on=False) + ax.imshow(test_A[k], cmap="binary", interpolation="nearest") + delta = 0.0 + ax.scatter( + test_bx[k, :, 0], + torch.full((test_bx.size(1),), delta), + color="black", + marker="s", + clip_on=False, + ) + ax.scatter( + torch.full((test_bx.size(1),), delta), + test_bx[k, :, 0], + color="black", + marker="s", + clip_on=False, + ) + ax.scatter( + test_tr[k, :, 0], + torch.full((test_tr.size(1),), delta), + color="black", + marker="^", + clip_on=False, + ) + ax.scatter( + torch.full((test_tr.size(1),), delta), + test_tr[k, :, 0], + color="black", + marker="^", + clip_on=False, + ) - fig.savefig(f'att1d_{label}test_A_{k:03d}.pdf', bbox_inches='tight') + fig.savefig(f"att1d_{label}test_A_{k:03d}.pdf", bbox_inches="tight") - plt.close('all') + plt.close("all") ###################################################################### diff --git a/causal-autoregression.py b/causal-autoregression.py index 7754265..0c931fb 100755 --- a/causal-autoregression.py +++ b/causal-autoregression.py @@ -18,46 +18,46 @@ from torch.nn import functional as F ###################################################################### -def save_images(x, filename, nrow = 12): - print(f'Writing {filename}') - torchvision.utils.save_image(x.narrow(0,0, min(48, x.size(0))), - filename, - nrow = nrow, pad_value=1.0) + +def save_images(x, filename, nrow=12): + print(f"Writing {filename}") + torchvision.utils.save_image( + x.narrow(0, 0, min(48, x.size(0))), filename, nrow=nrow, pad_value=1.0 + ) + ###################################################################### parser = argparse.ArgumentParser( - description = 'An implementation of a causal autoregression model', - formatter_class = argparse.ArgumentDefaultsHelpFormatter + description="An implementation of a causal autoregression model", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, ) -parser.add_argument('--data', - type = str, default = 'toy1d', - help = 'What data') +parser.add_argument("--data", type=str, default="toy1d", help="What data") -parser.add_argument('--seed', - type = int, default = 0, - help = 'Random seed (default 0, < 0 is no seeding)') +parser.add_argument( + "--seed", type=int, default=0, help="Random seed (default 0, < 0 is no seeding)" +) -parser.add_argument('--nb_epochs', - type = int, default = -1, - help = 'How many epochs') +parser.add_argument("--nb_epochs", type=int, default=-1, help="How many epochs") -parser.add_argument('--batch_size', - type = int, default = 100, - help = 'Batch size') +parser.add_argument("--batch_size", type=int, default=100, help="Batch size") -parser.add_argument('--learning_rate', - type = float, default = 1e-3, - help = 'Batch size') +parser.add_argument("--learning_rate", type=float, default=1e-3, help="Batch size") -parser.add_argument('--positional', - action='store_true', default = False, - help = 'Do we provide a positional encoding as input') +parser.add_argument( + "--positional", + action="store_true", + default=False, + help="Do we provide a positional encoding as input", +) -parser.add_argument('--dilation', - action='store_true', default = False, - help = 'Do we provide a positional encoding as input') +parser.add_argument( + "--dilation", + action="store_true", + default=False, + help="Do we provide a positional encoding as input", +) ###################################################################### @@ -67,32 +67,33 @@ if args.seed >= 0: torch.manual_seed(args.seed) if args.nb_epochs < 0: - if args.data == 'toy1d': + if args.data == "toy1d": args.nb_epochs = 100 - elif args.data == 'mnist': + elif args.data == "mnist": args.nb_epochs = 25 ###################################################################### if torch.cuda.is_available(): - print('Cuda is available') - device = torch.device('cuda') + print("Cuda is available") + device = torch.device("cuda") torch.backends.cudnn.benchmark = True else: - device = torch.device('cpu') + device = torch.device("cpu") ###################################################################### + class NetToy1d(nn.Module): - def __init__(self, nb_classes, ks = 2, nc = 32): + def __init__(self, nb_classes, ks=2, nc=32): super().__init__() self.pad = (ks - 1, 0) - self.conv0 = nn.Conv1d(1, nc, kernel_size = 1) - self.conv1 = nn.Conv1d(nc, nc, kernel_size = ks) - self.conv2 = nn.Conv1d(nc, nc, kernel_size = ks) - self.conv3 = nn.Conv1d(nc, nc, kernel_size = ks) - self.conv4 = nn.Conv1d(nc, nc, kernel_size = ks) - self.conv5 = nn.Conv1d(nc, nb_classes, kernel_size = 1) + self.conv0 = nn.Conv1d(1, nc, kernel_size=1) + self.conv1 = nn.Conv1d(nc, nc, kernel_size=ks) + self.conv2 = nn.Conv1d(nc, nc, kernel_size=ks) + self.conv3 = nn.Conv1d(nc, nc, kernel_size=ks) + self.conv4 = nn.Conv1d(nc, nc, kernel_size=ks) + self.conv5 = nn.Conv1d(nc, nb_classes, kernel_size=1) def forward(self, x): x = F.relu(self.conv0(F.pad(x, (1, -1)))) @@ -103,19 +104,20 @@ class NetToy1d(nn.Module): x = self.conv5(x) return x.permute(0, 2, 1).contiguous() + class NetToy1dWithDilation(nn.Module): - def __init__(self, nb_classes, ks = 2, nc = 32): + def __init__(self, nb_classes, ks=2, nc=32): super().__init__() - self.conv0 = nn.Conv1d(1, nc, kernel_size = 1) - self.pad1 = ((ks-1) * 2, 0) - self.conv1 = nn.Conv1d(nc, nc, kernel_size = ks, dilation = 2) - self.pad2 = ((ks-1) * 4, 0) - self.conv2 = nn.Conv1d(nc, nc, kernel_size = ks, dilation = 4) - self.pad3 = ((ks-1) * 8, 0) - self.conv3 = nn.Conv1d(nc, nc, kernel_size = ks, dilation = 8) - self.pad4 = ((ks-1) * 16, 0) - self.conv4 = nn.Conv1d(nc, nc, kernel_size = ks, dilation = 16) - self.conv5 = nn.Conv1d(nc, nb_classes, kernel_size = 1) + self.conv0 = nn.Conv1d(1, nc, kernel_size=1) + self.pad1 = ((ks - 1) * 2, 0) + self.conv1 = nn.Conv1d(nc, nc, kernel_size=ks, dilation=2) + self.pad2 = ((ks - 1) * 4, 0) + self.conv2 = nn.Conv1d(nc, nc, kernel_size=ks, dilation=4) + self.pad3 = ((ks - 1) * 8, 0) + self.conv3 = nn.Conv1d(nc, nc, kernel_size=ks, dilation=8) + self.pad4 = ((ks - 1) * 16, 0) + self.conv4 = nn.Conv1d(nc, nc, kernel_size=ks, dilation=16) + self.conv5 = nn.Conv1d(nc, nb_classes, kernel_size=1) def forward(self, x): x = F.relu(self.conv0(F.pad(x, (1, -1)))) @@ -126,21 +128,23 @@ class NetToy1dWithDilation(nn.Module): x = self.conv5(x) return x.permute(0, 2, 1).contiguous() + ###################################################################### + class PixelCNN(nn.Module): - def __init__(self, nb_classes, in_channels = 1, ks = 5): + def __init__(self, nb_classes, in_channels=1, ks=5): super().__init__() - self.hpad = (ks//2, ks//2, ks//2, 0) - self.vpad = (ks//2, 0, 0, 0) + self.hpad = (ks // 2, ks // 2, ks // 2, 0) + self.vpad = (ks // 2, 0, 0, 0) - self.conv1h = nn.Conv2d(in_channels, 32, kernel_size = (ks//2+1, ks)) - self.conv2h = nn.Conv2d(32, 64, kernel_size = (ks//2+1, ks)) - self.conv1v = nn.Conv2d(in_channels, 32, kernel_size = (1, ks//2+1)) - self.conv2v = nn.Conv2d(32, 64, kernel_size = (1, ks//2+1)) - self.final1 = nn.Conv2d(128, 128, kernel_size = 1) - self.final2 = nn.Conv2d(128, nb_classes, kernel_size = 1) + self.conv1h = nn.Conv2d(in_channels, 32, kernel_size=(ks // 2 + 1, ks)) + self.conv2h = nn.Conv2d(32, 64, kernel_size=(ks // 2 + 1, ks)) + self.conv1v = nn.Conv2d(in_channels, 32, kernel_size=(1, ks // 2 + 1)) + self.conv2v = nn.Conv2d(32, 64, kernel_size=(1, ks // 2 + 1)) + self.final1 = nn.Conv2d(128, 128, kernel_size=1) + self.final2 = nn.Conv2d(128, nb_classes, kernel_size=1) def forward(self, x): xh = F.pad(x, (0, 0, 1, -1)) @@ -154,8 +158,10 @@ class PixelCNN(nn.Module): return x.permute(0, 2, 3, 1).contiguous() + ###################################################################### + def positional_tensor(height, width): index_h = torch.arange(height).view(1, -1) m_h = (2 ** torch.arange(math.ceil(math.log2(height)))).view(-1, 1) @@ -169,26 +175,30 @@ def positional_tensor(height, width): return torch.cat((i_w, i_h), 1) + ###################################################################### str_experiment = args.data if args.positional: - str_experiment += '-positional' + str_experiment += "-positional" if args.dilation: - str_experiment += '-dilation' + str_experiment += "-dilation" + +log_file = open("causalar-" + str_experiment + "-train.log", "w") -log_file = open('causalar-' + str_experiment + '-train.log', 'w') def log_string(s): - s = time.strftime("%Y%m%d-%H:%M:%S", time.localtime()) + ' ' + s + s = time.strftime("%Y%m%d-%H:%M:%S", time.localtime()) + " " + s print(s) - log_file.write(s + '\n') + log_file.write(s + "\n") log_file.flush() + ###################################################################### + def generate_sequences(nb, len): nb_parts = 2 @@ -196,32 +206,33 @@ def generate_sequences(nb, len): x = torch.empty(nb, nb_parts).uniform_(-1, 1) x = x.view(nb, nb_parts, 1).expand(nb, nb_parts, len) - x = x * torch.linspace(0, len-1, len).view(1, -1) + len + x = x * torch.linspace(0, len - 1, len).view(1, -1) + len for n in range(nb): - a = torch.randperm(len - 2)[:nb_parts+1].sort()[0] + a = torch.randperm(len - 2)[: nb_parts + 1].sort()[0] a[0] = 0 a[a.size(0) - 1] = len for k in range(a.size(0) - 1): - r[n, a[k]:a[k+1]] = x[n, k, :a[k+1]-a[k]] + r[n, a[k] : a[k + 1]] = x[n, k, : a[k + 1] - a[k]] return r.round().long() + ###################################################################### -if args.data == 'toy1d': +if args.data == "toy1d": len = 32 train_input = generate_sequences(50000, len).to(device).unsqueeze(1) if args.dilation: - model = NetToy1dWithDilation(nb_classes = 2 * len).to(device) + model = NetToy1dWithDilation(nb_classes=2 * len).to(device) else: - model = NetToy1d(nb_classes = 2 * len).to(device) + model = NetToy1d(nb_classes=2 * len).to(device) -elif args.data == 'mnist': - train_set = torchvision.datasets.MNIST('./data/mnist/', train = True, download = True) +elif args.data == "mnist": + train_set = torchvision.datasets.MNIST("./data/mnist/", train=True, download=True) train_input = train_set.data.view(-1, 1, 28, 28).long().to(device) - model = PixelCNN(nb_classes = 256, in_channels = 1).to(device) + model = PixelCNN(nb_classes=256, in_channels=1).to(device) in_channels = train_input.size(1) if args.positional: @@ -229,40 +240,35 @@ elif args.data == 'mnist': positional_input = positional_tensor(height, width).float().to(device) in_channels += positional_input.size(1) - model = PixelCNN(nb_classes = 256, in_channels = in_channels).to(device) + model = PixelCNN(nb_classes=256, in_channels=in_channels).to(device) else: - raise ValueError('Unknown data ' + args.data) + raise ValueError("Unknown data " + args.data) ###################################################################### mean, std = train_input.float().mean(), train_input.float().std() nb_parameters = sum(t.numel() for t in model.parameters()) -log_string(f'nb_parameters {nb_parameters}') +log_string(f"nb_parameters {nb_parameters}") cross_entropy = nn.CrossEntropyLoss().to(device) -optimizer = torch.optim.Adam(model.parameters(), lr = args.learning_rate) +optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate) for e in range(args.nb_epochs): - nb_batches, acc_loss = 0, 0.0 for sequences in train_input.split(args.batch_size): - input = (sequences - mean)/std + input = (sequences - mean) / std if args.positional: input = torch.cat( - (input, positional_input.expand(input.size(0), -1, -1, -1)), - 1 + (input, positional_input.expand(input.size(0), -1, -1, -1)), 1 ) output = model(input) - loss = cross_entropy( - output.view(-1, output.size(-1)), - sequences.view(-1) - ) + loss = cross_entropy(output.view(-1, output.size(-1)), sequences.view(-1)) optimizer.zero_grad() loss.backward() @@ -271,7 +277,7 @@ for e in range(args.nb_epochs): nb_batches += 1 acc_loss += loss.item() - log_string(f'{e} {acc_loss / nb_batches} {math.exp(acc_loss / nb_batches)}') + log_string(f"{e} {acc_loss / nb_batches} {math.exp(acc_loss / nb_batches)}") sys.stdout.flush() @@ -284,36 +290,36 @@ flat = generated.view(generated.size(0), -1) for t in range(flat.size(1)): input = (generated.float() - mean) / std if args.positional: - input = torch.cat((input, positional_input.expand(input.size(0), -1, -1, -1)), 1) + input = torch.cat( + (input, positional_input.expand(input.size(0), -1, -1, -1)), 1 + ) output = model(input) logits = output.view(flat.size() + (-1,))[:, t] - dist = torch.distributions.categorical.Categorical(logits = logits) + dist = torch.distributions.categorical.Categorical(logits=logits) flat[:, t] = dist.sample() ###################################################################### -if args.data == 'toy1d': - - with open('causalar-' + str_experiment + '-train.dat', 'w') as file: +if args.data == "toy1d": + with open("causalar-" + str_experiment + "-train.dat", "w") as file: for j in range(train_input.size(2)): - file.write(f'{j}') + file.write(f"{j}") for i in range(min(train_input.size(0), 25)): - file.write(f' {train_input[i, 0, j]}') - file.write('\n') + file.write(f" {train_input[i, 0, j]}") + file.write("\n") - with open('causalar-' + str_experiment + '-generated.dat', 'w') as file: + with open("causalar-" + str_experiment + "-generated.dat", "w") as file: for j in range(generated.size(2)): - file.write(f'{j}') + file.write(f"{j}") for i in range(generated.size(0)): - file.write(f' {generated[i, 0, j]}') - file.write('\n') - -elif args.data == 'mnist': + file.write(f" {generated[i, 0, j]}") + file.write("\n") - img_train = 1 - train_input[:generated.size(0)].float() / 255 +elif args.data == "mnist": + img_train = 1 - train_input[: generated.size(0)].float() / 255 img_generated = 1 - generated.float() / 255 - save_images(img_train, 'causalar-' + str_experiment + '-train.png', nrow = 12) - save_images(img_generated, 'causalar-' + str_experiment + '-generated.png', nrow = 12) + save_images(img_train, "causalar-" + str_experiment + "-train.png", nrow=12) + save_images(img_generated, "causalar-" + str_experiment + "-generated.png", nrow=12) ###################################################################### diff --git a/confidence.py b/confidence.py index 4530fbc..a586a3d 100755 --- a/confidence.py +++ b/confidence.py @@ -30,19 +30,24 @@ y = y.view(-1, 1) nh = 400 -model = nn.Sequential(nn.Linear(1, nh), nn.ReLU(), - nn.Dropout(0.25), - nn.Linear(nh, nh), nn.ReLU(), - nn.Dropout(0.25), - nn.Linear(nh, 1)) +model = nn.Sequential( + nn.Linear(1, nh), + nn.ReLU(), + nn.Dropout(0.25), + nn.Linear(nh, nh), + nn.ReLU(), + nn.Dropout(0.25), + nn.Linear(nh, 1), +) model.train(True) criterion = nn.MSELoss() -optimizer = torch.optim.Adam(model.parameters(), lr = 1e-4) +optimizer = torch.optim.Adam(model.parameters(), lr=1e-4) for k in range(10000): loss = criterion(model(x), y) - if (k+1)%100 == 0: print(k+1, loss.item()) + if (k + 1) % 100 == 0: + print(k + 1, loss.item()) optimizer.zero_grad() loss.backward() optimizer.step() @@ -59,8 +64,13 @@ v = model(v).reshape(101, -1) mean = v.mean(1) std = v.std(1) -ax.fill_between(u.numpy(), (mean-std).detach().numpy(), (mean+std).detach().numpy(), color = '#e0e0e0') -ax.plot(u.numpy(), mean.detach().numpy(), color = 'red') +ax.fill_between( + u.numpy(), + (mean - std).detach().numpy(), + (mean + std).detach().numpy(), + color="#e0e0e0", +) +ax.plot(u.numpy(), mean.detach().numpy(), color="red") ax.scatter(x.numpy(), y.numpy()) plt.show() diff --git a/conv_chain.py b/conv_chain.py index 3077874..a1d9af0 100755 --- a/conv_chain.py +++ b/conv_chain.py @@ -7,44 +7,49 @@ ###################################################################### + def conv_chain(input_size, output_size, remain_depth, cond): if remain_depth == 0: if input_size == output_size: - return [ [ ] ] + return [[]] else: - return [ ] + return [] else: - r = [ ] + r = [] for kernel_size in range(1, input_size + 1): for stride in range(1, input_size): if cond(remain_depth, kernel_size, stride): n = (input_size - kernel_size) // stride + 1 - if n >= output_size and (n - 1) * stride + kernel_size == input_size: + if ( + n >= output_size + and (n - 1) * stride + kernel_size == input_size + ): q = conv_chain(n, output_size, remain_depth - 1, cond) - r += [ [ (kernel_size, stride) ] + u for u in q ] + r += [[(kernel_size, stride)] + u for u in q] return r + ###################################################################### if __name__ == "__main__": - import torch from torch import nn # Example c = conv_chain( - input_size = 64, output_size = 8, - remain_depth = 5, + input_size=64, + output_size=8, + remain_depth=5, # We want kernels smaller than 4, strides smaller than the # kernels, and strides of 1 except in the two last layers - cond = lambda d, k, s: k <= 4 and s <= k and (s == 1 or d <= 2) + cond=lambda d, k, s: k <= 4 and s <= k and (s == 1 or d <= 2), ) x = torch.rand(1, 1, 64) for m in c: - model = nn.Sequential(*[ nn.Conv1d(1, 1, l[0], l[1]) for l in m ]) + model = nn.Sequential(*[nn.Conv1d(1, 1, l[0], l[1]) for l in m]) print(model) print(x.size(), model(x).size()) diff --git a/ddpol.py b/ddpol.py index 645f47c..1975ab2 100755 --- a/ddpol.py +++ b/ddpol.py @@ -12,23 +12,21 @@ import torch ###################################################################### -parser = argparse.ArgumentParser(description='Example of double descent with polynomial regression.') +parser = argparse.ArgumentParser( + description="Example of double descent with polynomial regression." +) -parser.add_argument('--D-max', - type = int, default = 16) +parser.add_argument("--D-max", type=int, default=16) -parser.add_argument('--nb-runs', - type = int, default = 250) +parser.add_argument("--nb-runs", type=int, default=250) -parser.add_argument('--nb-train-samples', - type = int, default = 8) +parser.add_argument("--nb-train-samples", type=int, default=8) -parser.add_argument('--train-noise-std', - type = float, default = 0.) +parser.add_argument("--train-noise-std", type=float, default=0.0) -parser.add_argument('--seed', - type = int, default = 0, - help = 'Random seed (default 0, < 0 is no seeding)') +parser.add_argument( + "--seed", type=int, default=0, help="Random seed (default 0, < 0 is no seeding)" +) args = parser.parse_args() @@ -37,80 +35,101 @@ if args.seed >= 0: ###################################################################### + def pol_value(alpha, x): x_pow = x.view(-1, 1) ** torch.arange(alpha.size(0)).view(1, -1) return x_pow @ alpha -def fit_alpha(x, y, D, a = 0, b = 1, rho = 1e-12): + +def fit_alpha(x, y, D, a=0, b=1, rho=1e-12): M = x.view(-1, 1) ** torch.arange(D + 1).view(1, -1) B = y if D >= 2: - q = torch.arange(2, D + 1, dtype = x.dtype).view(1, -1) - r = q.view(-1, 1) + q = torch.arange(2, D + 1, dtype=x.dtype).view(1, -1) + r = q.view(-1, 1) beta = x.new_zeros(D + 1, D + 1) - beta[2:, 2:] = (q-1) * q * (r-1) * r * (b**(q+r-3) - a**(q+r-3))/(q+r-3) + beta[2:, 2:] = ( + (q - 1) + * q + * (r - 1) + * r + * (b ** (q + r - 3) - a ** (q + r - 3)) + / (q + r - 3) + ) W = torch.linalg.eig(beta) l, U = W.eigenvalues.real, W.eigenvectors.real - Q = U @ torch.diag(l.clamp(min = 0) ** 0.5) # clamp deals with ~0 negative values + Q = U @ torch.diag(l.clamp(min=0) ** 0.5) # clamp deals with ~0 negative values B = torch.cat((B, y.new_zeros(Q.size(0))), 0) M = torch.cat((M, math.sqrt(rho) * Q.t()), 0) - return torch.linalg.lstsq(M, B).solution[:D+1] + return torch.linalg.lstsq(M, B).solution[: D + 1] + ###################################################################### # The "ground truth" + def phi(x): - return torch.abs(torch.abs(x - 0.4) - 0.2) + x/2 - 0.1 + return torch.abs(torch.abs(x - 0.4) - 0.2) + x / 2 - 0.1 + ###################################################################### + def compute_mse(nb_train_samples): mse_train = torch.zeros(args.nb_runs, args.D_max + 1) mse_test = torch.zeros(args.nb_runs, args.D_max + 1) for k in range(args.nb_runs): - x_train = torch.rand(nb_train_samples, dtype = torch.float64) + x_train = torch.rand(nb_train_samples, dtype=torch.float64) y_train = phi(x_train) if args.train_noise_std > 0: - y_train = y_train + torch.empty_like(y_train).normal_(0, args.train_noise_std) - x_test = torch.linspace(0, 1, 100, dtype = x_train.dtype) + y_train = y_train + torch.empty_like(y_train).normal_( + 0, args.train_noise_std + ) + x_test = torch.linspace(0, 1, 100, dtype=x_train.dtype) y_test = phi(x_test) for D in range(args.D_max + 1): alpha = fit_alpha(x_train, y_train, D) - mse_train[k, D] = ((pol_value(alpha, x_train) - y_train)**2).mean() - mse_test[k, D] = ((pol_value(alpha, x_test) - y_test)**2).mean() + mse_train[k, D] = ((pol_value(alpha, x_train) - y_train) ** 2).mean() + mse_test[k, D] = ((pol_value(alpha, x_test) - y_test) ** 2).mean() return mse_train.median(0).values, mse_test.median(0).values + ###################################################################### # Plot the MSE vs. degree curves fig = plt.figure() ax = fig.add_subplot(1, 1, 1) -ax.set_yscale('log') +ax.set_yscale("log") ax.set_ylim(1e-5, 1) -ax.set_xlabel('Polynomial degree', labelpad = 10) -ax.set_ylabel('MSE', labelpad = 10) +ax.set_xlabel("Polynomial degree", labelpad=10) +ax.set_ylabel("MSE", labelpad=10) -ax.axvline(x = args.nb_train_samples - 1, - color = 'gray', linewidth = 0.5, linestyle = '--') +ax.axvline(x=args.nb_train_samples - 1, color="gray", linewidth=0.5, linestyle="--") -ax.text(args.nb_train_samples - 1.2, 1e-4, 'nb. params = nb. samples', - fontsize = 10, color = 'gray', - rotation = 90, rotation_mode='anchor') +ax.text( + args.nb_train_samples - 1.2, + 1e-4, + "nb. params = nb. samples", + fontsize=10, + color="gray", + rotation=90, + rotation_mode="anchor", +) mse_train, mse_test = compute_mse(args.nb_train_samples) -ax.plot(torch.arange(args.D_max + 1), mse_train, color = 'blue', label = 'Train') -ax.plot(torch.arange(args.D_max + 1), mse_test, color = 'red', label = 'Test') +ax.plot(torch.arange(args.D_max + 1), mse_train, color="blue", label="Train") +ax.plot(torch.arange(args.D_max + 1), mse_test, color="red", label="Test") -ax.legend(frameon = False) +ax.legend(frameon=False) -fig.savefig('dd-mse.pdf', bbox_inches='tight') +fig.savefig("dd-mse.pdf", bbox_inches="tight") plt.close(fig) @@ -120,54 +139,66 @@ plt.close(fig) fig = plt.figure() ax = fig.add_subplot(1, 1, 1) -ax.set_yscale('log') +ax.set_yscale("log") ax.set_ylim(1e-5, 1) -ax.set_xlabel('Polynomial degree', labelpad = 10) -ax.set_ylabel('MSE', labelpad = 10) +ax.set_xlabel("Polynomial degree", labelpad=10) +ax.set_ylabel("MSE", labelpad=10) nb_train_samples_min = args.nb_train_samples - 4 nb_train_samples_max = args.nb_train_samples for nb_train_samples in range(nb_train_samples_min, nb_train_samples_max + 1, 2): mse_train, mse_test = compute_mse(nb_train_samples) - e = float(nb_train_samples - nb_train_samples_min) / float(nb_train_samples_max - nb_train_samples_min) + e = float(nb_train_samples - nb_train_samples_min) / float( + nb_train_samples_max - nb_train_samples_min + ) e = 0.15 + 0.7 * e - ax.plot(torch.arange(args.D_max + 1), mse_train, color = (e, e, 1.0), label = f'Train N={nb_train_samples}') - ax.plot(torch.arange(args.D_max + 1), mse_test, color = (1.0, e, e), label = f'Test N={nb_train_samples}') - -ax.legend(frameon = False) - -fig.savefig('dd-multi-mse.pdf', bbox_inches='tight') + ax.plot( + torch.arange(args.D_max + 1), + mse_train, + color=(e, e, 1.0), + label=f"Train N={nb_train_samples}", + ) + ax.plot( + torch.arange(args.D_max + 1), + mse_test, + color=(1.0, e, e), + label=f"Test N={nb_train_samples}", + ) + +ax.legend(frameon=False) + +fig.savefig("dd-multi-mse.pdf", bbox_inches="tight") plt.close(fig) ###################################################################### # Plot some examples of train / test -torch.manual_seed(9) # I picked that for pretty +torch.manual_seed(9) # I picked that for pretty -x_train = torch.rand(args.nb_train_samples, dtype = torch.float64) +x_train = torch.rand(args.nb_train_samples, dtype=torch.float64) y_train = phi(x_train) if args.train_noise_std > 0: y_train = y_train + torch.empty_like(y_train).normal_(0, args.train_noise_std) -x_test = torch.linspace(0, 1, 100, dtype = x_train.dtype) +x_test = torch.linspace(0, 1, 100, dtype=x_train.dtype) y_test = phi(x_test) for D in range(args.D_max + 1): fig = plt.figure() ax = fig.add_subplot(1, 1, 1) - ax.set_title(f'Degree {D}') + ax.set_title(f"Degree {D}") ax.set_ylim(-0.1, 1.1) - ax.plot(x_test, y_test, color = 'black', label = 'Test values') - ax.scatter(x_train, y_train, color = 'blue', label = 'Train samples') + ax.plot(x_test, y_test, color="black", label="Test values") + ax.scatter(x_train, y_train, color="blue", label="Train samples") alpha = fit_alpha(x_train, y_train, D) - ax.plot(x_test, pol_value(alpha, x_test), color = 'red', label = 'Fitted polynomial') + ax.plot(x_test, pol_value(alpha, x_test), color="red", label="Fitted polynomial") - ax.legend(frameon = False) + ax.legend(frameon=False) - fig.savefig(f'dd-example-{D:02d}.pdf', bbox_inches='tight') + fig.savefig(f"dd-example-{D:02d}.pdf", bbox_inches="tight") plt.close(fig) diff --git a/denoising-ae-field.py b/denoising-ae-field.py index f96c23a..3ef0c80 100755 --- a/denoising-ae-field.py +++ b/denoising-ae-field.py @@ -13,34 +13,35 @@ from torch import nn ###################################################################### + def data_rectangle(nb): x = torch.rand(nb, 1) - 0.5 y = torch.rand(nb, 1) * 2 - 1 data = torch.cat((y, x), 1) alpha = math.pi / 8 data = data @ torch.tensor( - [ - [ math.cos(alpha), math.sin(alpha)], - [-math.sin(alpha), math.cos(alpha)] - ] + [[math.cos(alpha), math.sin(alpha)], [-math.sin(alpha), math.cos(alpha)]] ) - return data, 'rectangle' + return data, "rectangle" + def data_zigzag(nb): a = torch.empty(nb).uniform_(0, 1).view(-1, 1) # zigzag - x = 0.4 * ((a-0.5) * 5 * math.pi).cos() + x = 0.4 * ((a - 0.5) * 5 * math.pi).cos() y = a * 2.5 - 1.25 data = torch.cat((y, x), 1) - data = data @ torch.tensor([[1., -1.], [1., 1.]]) - return data, 'zigzag' + data = data @ torch.tensor([[1.0, -1.0], [1.0, 1.0]]) + 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, 'spiral' + return data, "spiral" + def data_penta(nb): a = (torch.randint(5, (nb,)).float() / 5 * 2 * math.pi).view(-1, 1) @@ -48,19 +49,17 @@ 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, 'penta' + return data, "penta" + ###################################################################### + def train_model(data): - model = nn.Sequential( - nn.Linear(2, 100), - nn.ReLU(), - nn.Linear(100, 2) - ) + 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) + optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) criterion = nn.MSELoss() for e in range(nb_epochs): @@ -73,16 +72,19 @@ def train_model(data): optimizer.zero_grad() loss.backward() optimizer.step() - if (e+1)%100 == 0: print(e+1, acc_loss) + if (e + 1) % 100 == 0: + print(e + 1, acc_loss) return model + ###################################################################### + def save_image(data_name, model, data): 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) + 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 the grid closer than @@ -95,30 +97,35 @@ def save_image(data_name, model, data): fig = plt.figure() ax = fig.add_subplot(1, 1, 1) - ax.axis('off') + 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 + 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' + data[:, 0].numpy(), data[:, 1].numpy(), s=1, color="tab:blue" ) - filename = f'denoising_field_{data_name}.pdf' - print(f'Saving {filename}') - fig.savefig(filename, 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_rectangle, data_zigzag, data_spiral, data_penta ]: +for data_source in [data_rectangle, data_zigzag, data_spiral, data_penta]: data, data_name = data_source(1000) data = data - data.mean(0) model = train_model(data) diff --git a/elbo.py b/elbo.py index 6af4a77..dbea3b5 100755 --- a/elbo.py +++ b/elbo.py @@ -7,8 +7,10 @@ import torch + def D_KL(a, b): - return - a @ (b / a).log() + return -a @ (b / a).log() + # p(X = x, Z = z) = p[x, z] @@ -19,12 +21,12 @@ q_XZ /= q_XZ.sum() p_X = p_XZ.sum(1) p_Z = p_XZ.sum(0) -p_X_given_Z = p_XZ / p_XZ.sum(0, keepdim = True) -p_Z_given_X = p_XZ / p_XZ.sum(1, keepdim = True) +p_X_given_Z = p_XZ / p_XZ.sum(0, keepdim=True) +p_Z_given_X = p_XZ / p_XZ.sum(1, keepdim=True) -#q_X_given_Z = q_XZ / q_XZ.sum(0, keepdim = True) -q_Z_given_X = q_XZ / q_XZ.sum(1, keepdim = True) +# q_X_given_Z = q_XZ / q_XZ.sum(0, keepdim = True) +q_Z_given_X = q_XZ / q_XZ.sum(1, keepdim=True) for x in range(p_XZ.size(0)): - elbo = q_Z_given_X[x, :] @ ( p_X_given_Z[x, :] / q_Z_given_X[x, :] * p_Z).log() + elbo = q_Z_given_X[x, :] @ (p_X_given_Z[x, :] / q_Z_given_X[x, :] * p_Z).log() print(p_X[x].log(), elbo + D_KL(q_Z_given_X[x, :], p_Z_given_X[x, :])) diff --git a/flatparam.py b/flatparam.py index 57a8720..0b61cf1 100755 --- a/flatparam.py +++ b/flatparam.py @@ -5,12 +5,13 @@ from torch import nn ###################################################################### -def _flatparam(model, whole, already = [], offset = 0): + +def _flatparam(model, whole, already=[], offset=0): for v in model._parameters: p = model._parameters[v] e = p.numel() s = p.size() - model._parameters[v] = whole[offset:offset+e].view(s) + model._parameters[v] = whole[offset : offset + e].view(s) with torch.no_grad(): model._parameters[v].copy_(p) offset += e @@ -20,44 +21,42 @@ def _flatparam(model, whole, already = [], offset = 0): offset = _flatparam(m, whole, already, offset) return offset + def flatparam(model): n = sum(p.numel() for p in model.parameters()) - whole = next(model.parameters()).new(n) # Get same device and dtype + whole = next(model.parameters()).new(n) # Get same device and dtype whole.requires_grad_() _flatparam(model, whole) - model.parameters = lambda: iter([ whole ]) + model.parameters = lambda: iter([whole]) + ###################################################################### model = nn.Sequential( nn.Linear(2, 4), nn.ReLU(), - nn.Sequential( - nn.Linear(4, 4), - nn.ReLU(), - nn.Linear(4, 2) - ) + nn.Sequential(nn.Linear(4, 4), nn.ReLU(), nn.Linear(4, 2)), ) ###################################################################### -print('Before:') +print("Before:") for p in model.parameters(): print(p.size(), p.storage().size()) flatparam(model) -print('After:') +print("After:") for p in model.parameters(): print(p.size(), p.storage().size()) ###################################################################### -print('Check:') +print("Check:") input = torch.rand(100, 2) targets = torch.rand(100, 2) -optimizer = torch.optim.SGD(model.parameters(), lr = 1e-2) +optimizer = torch.optim.SGD(model.parameters(), lr=1e-2) mse = nn.MSELoss() for e in range(10): diff --git a/gpt-test.py b/gpt-test.py index 0967043..ddd7dcf 100755 --- a/gpt-test.py +++ b/gpt-test.py @@ -22,9 +22,10 @@ from transformers import GPT2Tokenizer, GPT2LMHeadModel ###################################################################### -def complete(model, tokenizer, - primer, - nb_sentences = 1, nb_token_max = 100, temperature = None): + +def complete( + model, tokenizer, primer, nb_sentences=1, nb_token_max=100, temperature=None +): nt, ns = 0, 0 tokens = tokenizer.encode(primer) primer_len = len(tokens) @@ -33,31 +34,41 @@ def complete(model, tokenizer, if temperature is None: next_token = torch.argmax(outputs[0, -1]) else: - dist = torch.distributions.Categorical(logits = outputs[0, -1] / temperature) + dist = torch.distributions.Categorical(logits=outputs[0, -1] / temperature) next_token = dist.sample((1,)).item() tokens.append(next_token) nt += 1 - if tokenizer.decode([next_token]) == '.': ns += 1 + if tokenizer.decode([next_token]) == ".": + ns += 1 if ns == nb_sentences or nt == nb_token_max: - return '<' + tokenizer.decode(tokens[:primer_len]) + '>' + \ - tokenizer.decode(tokens[primer_len:]) + return ( + "<" + + tokenizer.decode(tokens[:primer_len]) + + ">" + + tokenizer.decode(tokens[primer_len:]) + ) + ###################################################################### -#model_name = 'gpt2' -#model_name = 'gpt2-large' -model_name = 'gpt2-xl' +# model_name = 'gpt2' +# model_name = 'gpt2-large' +model_name = "gpt2-xl" tokenizer = GPT2Tokenizer.from_pretrained(model_name) model = GPT2LMHeadModel.from_pretrained(model_name) model.eval() -print(f'Using {model_name} ({int(sum(p.numel() for p in model.parameters())/(1e6))}M parameters)') +print( + f"Using {model_name} ({int(sum(p.numel() for p in model.parameters())/(1e6))}M parameters)" +) print( - complete(model, tokenizer, - 'The object was blue all over, but also green all over, it was a', + complete( + model, + tokenizer, + "The object was blue all over, but also green all over, it was a", ) ) diff --git a/hallu.py b/hallu.py index de25188..9d2706d 100755 --- a/hallu.py +++ b/hallu.py @@ -12,10 +12,11 @@ import PIL, torch, torchvision from torch.nn import functional as F + class MultiScaleEdgeEnergy(torch.nn.Module): def __init__(self): super().__init__() - k = torch.exp(- torch.tensor([[-2., -1., 0., 1., 2.]])**2 / 2) + k = torch.exp(-torch.tensor([[-2.0, -1.0, 0.0, 1.0, 2.0]]) ** 2 / 2) k = (k.t() @ k).view(1, 1, 5, 5) self.gaussian_5x5 = torch.nn.Parameter(k / k.sum()).requires_grad_(False) @@ -23,19 +24,20 @@ class MultiScaleEdgeEnergy(torch.nn.Module): u = x.view(-1, 1, x.size(2), x.size(3)) result = 0.0 while min(u.size(2), u.size(3)) > 5: - blurry = F.conv2d(u, self.gaussian_5x5, padding = 2) + blurry = F.conv2d(u, self.gaussian_5x5, padding=2) result += (u - blurry).view(u.size(0), -1).pow(2).sum(1) - u = F.avg_pool2d(u, kernel_size = 2, padding = 1) + u = F.avg_pool2d(u, kernel_size=2, padding=1) return result.view(x.size(0), -1).sum(1) -img = torchvision.transforms.ToTensor()(PIL.Image.open('blacklab.jpg')) + +img = torchvision.transforms.ToTensor()(PIL.Image.open("blacklab.jpg")) img = img.view((1,) + img.size()) ref_input = 0.5 + 0.5 * (img - img.mean()) / img.std() mse_loss = torch.nn.MSELoss() edge_energy = MultiScaleEdgeEnergy() -layers = torchvision.models.vgg16(pretrained = True).features +layers = torchvision.models.vgg16(pretrained=True).features layers.eval() if torch.cuda.is_available(): @@ -43,13 +45,13 @@ if torch.cuda.is_available(): ref_input = ref_input.cuda() layers.cuda() -for l in [ 5, 7, 12, 17, 21, 28 ]: +for l in [5, 7, 12, 17, 21, 28]: model = torch.nn.Sequential(layers[:l]) ref_output = model(ref_input).detach() for n in range(5): input = torch.empty_like(ref_input).uniform_(-0.01, 0.01).requires_grad_() - optimizer = torch.optim.Adam( [ input ], lr = 1e-2) + optimizer = torch.optim.Adam([input], lr=1e-2) for k in range(1000): output = model(input) loss = mse_loss(output, ref_output) + 1e-3 * edge_energy(input) @@ -58,7 +60,7 @@ for l in [ 5, 7, 12, 17, 21, 28 ]: optimizer.step() img = 0.5 + 0.2 * (input - input.mean()) / input.std() - result_name = 'hallu-l%02d-n%02d.png' % (l, n) + result_name = "hallu-l%02d-n%02d.png" % (l, n) torchvision.utils.save_image(img, result_name) - print('Wrote ' + result_name) + print("Wrote " + result_name) diff --git a/lazy_linear.py b/lazy_linear.py index 97530ef..c49f0d0 100755 --- a/lazy_linear.py +++ b/lazy_linear.py @@ -9,9 +9,9 @@ from torch import nn, Tensor ###################################################################### -class LazyLinear(nn.Module): - def __init__(self, out_dim, bias = True): +class LazyLinear(nn.Module): + def __init__(self, out_dim, bias=True): super().__init__() self.out_dim = out_dim self.bias = bias @@ -24,22 +24,25 @@ class LazyLinear(nn.Module): if self.training: self.core = nn.Linear(x.size(1), self.out_dim, self.bias) else: - raise RuntimeError('Undefined LazyLinear core in inference mode.') + raise RuntimeError("Undefined LazyLinear core in inference mode.") return self.core(x) - def named_parameters(self, memo=None, prefix=''): - assert self.core is not None, 'Parameters not yet defined' + def named_parameters(self, memo=None, prefix=""): + assert self.core is not None, "Parameters not yet defined" return super().named_parameters(memo, prefix) + ###################################################################### if __name__ == "__main__": - model = nn.Sequential(nn.Conv2d(3, 8, kernel_size = 5), - nn.ReLU(inplace = True), - LazyLinear(128), - nn.ReLU(inplace = True), - nn.Linear(128, 10)) + model = nn.Sequential( + nn.Conv2d(3, 8, kernel_size=5), + nn.ReLU(inplace=True), + LazyLinear(128), + nn.ReLU(inplace=True), + nn.Linear(128, 10), + ) # model.eval() @@ -49,4 +52,3 @@ if __name__ == "__main__": for n, x in model.named_parameters(): print(n, x.size()) - diff --git a/mandelbrot.py b/mandelbrot.py index fa522eb..c284eb0 100755 --- a/mandelbrot.py +++ b/mandelbrot.py @@ -13,4 +13,4 @@ zi = torch.zeros(n, n) for k in range(100): zr, zi = zr**2 - zi**2 + cr, 2 * zr * zi + ci -torchvision.utils.save_image(1-(1-zr**2 + zi**2).sign(), 'mandelbrot.png') +torchvision.utils.save_image(1 - (1 - zr**2 + zi**2).sign(), "mandelbrot.png") diff --git a/mi_estimator.py b/mi_estimator.py index 47381ef..1a167fe 100755 --- a/mi_estimator.py +++ b/mi_estimator.py @@ -17,14 +17,14 @@ import torch.nn.functional as F if torch.cuda.is_available(): torch.backends.cudnn.benchmark = True - device = torch.device('cuda') + device = torch.device("cuda") else: - device = torch.device('cpu') + device = torch.device("cpu") ###################################################################### parser = argparse.ArgumentParser( - description = '''An implementation of a Mutual Information estimator with a deep model + description="""An implementation of a Mutual Information estimator with a deep model Three different toy data-sets are implemented, each consists of pairs of samples, that may be from different spaces: @@ -39,41 +39,43 @@ parser = argparse.ArgumentParser( (3) Two 1d sequences, the first with a single peak, the second with two peaks, and the height of the peak in the first is the difference of timing of the peaks in the second. The "true" MI is - the log of the number of possible peak heights.''', - - formatter_class = argparse.ArgumentDefaultsHelpFormatter + the log of the number of possible peak heights.""", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, ) -parser.add_argument('--data', - type = str, default = 'image_pair', - help = 'What data: image_pair, image_values_pair, sequence_pair') +parser.add_argument( + "--data", + type=str, + default="image_pair", + help="What data: image_pair, image_values_pair, sequence_pair", +) -parser.add_argument('--seed', - type = int, default = 0, - help = 'Random seed (default 0, < 0 is no seeding)') +parser.add_argument( + "--seed", type=int, default=0, help="Random seed (default 0, < 0 is no seeding)" +) -parser.add_argument('--mnist_classes', - type = str, default = '0, 1, 3, 5, 6, 7, 8, 9', - help = 'What MNIST classes to use') +parser.add_argument( + "--mnist_classes", + type=str, + default="0, 1, 3, 5, 6, 7, 8, 9", + help="What MNIST classes to use", +) -parser.add_argument('--nb_classes', - type = int, default = 2, - help = 'How many classes for sequences') +parser.add_argument( + "--nb_classes", type=int, default=2, help="How many classes for sequences" +) -parser.add_argument('--nb_epochs', - type = int, default = 50, - help = 'How many epochs') +parser.add_argument("--nb_epochs", type=int, default=50, help="How many epochs") -parser.add_argument('--batch_size', - type = int, default = 100, - help = 'Batch size') +parser.add_argument("--batch_size", type=int, default=100, help="Batch size") -parser.add_argument('--learning_rate', - type = float, default = 1e-3, - help = 'Batch size') +parser.add_argument("--learning_rate", type=float, default=1e-3, help="Batch size") -parser.add_argument('--independent', action = 'store_true', - help = 'Should the pair components be independent') +parser.add_argument( + "--independent", + action="store_true", + help="Should the pair components be independent", +) ###################################################################### @@ -82,26 +84,29 @@ args = parser.parse_args() if args.seed >= 0: torch.manual_seed(args.seed) -used_MNIST_classes = torch.tensor(eval('[' + args.mnist_classes + ']'), device = device) +used_MNIST_classes = torch.tensor(eval("[" + args.mnist_classes + "]"), device=device) ###################################################################### + def entropy(target): probas = [] for k in range(target.max() + 1): n = (target == k).sum().item() - if n > 0: probas.append(n) + if n > 0: + probas.append(n) probas = torch.tensor(probas).float() probas /= probas.sum() - return - (probas * probas.log()).sum().item() + return -(probas * probas.log()).sum().item() + ###################################################################### -train_set = torchvision.datasets.MNIST('./data/mnist/', train = True, download = True) -train_input = train_set.train_data.view(-1, 1, 28, 28).to(device).float() +train_set = torchvision.datasets.MNIST("./data/mnist/", train=True, download=True) +train_input = train_set.train_data.view(-1, 1, 28, 28).to(device).float() train_target = train_set.train_labels.to(device) -test_set = torchvision.datasets.MNIST('./data/mnist/', train = False, download = True) +test_set = torchvision.datasets.MNIST("./data/mnist/", train=False, download=True) test_input = test_set.test_data.view(-1, 1, 28, 28).to(device).float() test_target = test_set.test_labels.to(device) @@ -115,7 +120,8 @@ test_input.sub_(mu).div_(std) # half of the samples, with a[i] and b[i] of same class for any i, and # c is a 1d long tensor real classes -def create_image_pairs(train = False): + +def create_image_pairs(train=False): ua, ub, uc = [], [], [] if train: @@ -124,11 +130,12 @@ def create_image_pairs(train = False): input, target = test_input, test_target for i in used_MNIST_classes: - used_indices = torch.arange(input.size(0), device = target.device)\ - .masked_select(target == i.item()) + used_indices = torch.arange(input.size(0), device=target.device).masked_select( + target == i.item() + ) x = input[used_indices] x = x[torch.randperm(x.size(0))] - hs = x.size(0)//2 + hs = x.size(0) // 2 ua.append(x.narrow(0, 0, hs)) ub.append(x.narrow(0, hs, hs)) uc.append(target[used_indices]) @@ -145,6 +152,7 @@ def create_image_pairs(train = False): return a, b, c + ###################################################################### # Returns a triplet a, b, c where a are the standard MNIST images, c @@ -153,7 +161,8 @@ def create_image_pairs(train = False): # b[n, 0] ~ Uniform(0, 10) # b[n, 1] ~ b[n, 0] + Uniform(0, 0.5) + c[n] -def create_image_values_pairs(train = False): + +def create_image_values_pairs(train=False): ua, ub = [], [] if train: @@ -161,10 +170,12 @@ def create_image_values_pairs(train = False): else: input, target = test_input, test_target - m = torch.zeros(used_MNIST_classes.max() + 1, dtype = torch.uint8, device = target.device) + m = torch.zeros( + used_MNIST_classes.max() + 1, dtype=torch.uint8, device=target.device + ) m[used_MNIST_classes] = 1 m = m[target] - used_indices = torch.arange(input.size(0), device = target.device).masked_select(m) + used_indices = torch.arange(input.size(0), device=target.device).masked_select(m) input = input[used_indices].contiguous() target = target[used_indices].contiguous() @@ -177,42 +188,46 @@ def create_image_values_pairs(train = False): b[:, 1].uniform_(0.0, 0.5) if args.independent: - b[:, 1] += b[:, 0] + \ - used_MNIST_classes[torch.randint(len(used_MNIST_classes), target.size())] + b[:, 1] += ( + b[:, 0] + + used_MNIST_classes[torch.randint(len(used_MNIST_classes), target.size())] + ) else: b[:, 1] += b[:, 0] + target.float() return a, b, c + ###################################################################### # -def create_sequences_pairs(train = False): + +def create_sequences_pairs(train=False): nb, length = 10000, 1024 noise_level = 2e-2 - ha = torch.randint(args.nb_classes, (nb, ), device = device) + 1 + ha = torch.randint(args.nb_classes, (nb,), device=device) + 1 if args.independent: - hb = torch.randint(args.nb_classes, (nb, ), device = device) + hb = torch.randint(args.nb_classes, (nb,), device=device) else: hb = ha - pos = torch.empty(nb, device = device).uniform_(0.0, 0.9) - a = torch.linspace(0, 1, length, device = device).view(1, -1).expand(nb, -1) + pos = torch.empty(nb, device=device).uniform_(0.0, 0.9) + a = torch.linspace(0, 1, length, device=device).view(1, -1).expand(nb, -1) a = a - pos.view(nb, 1) a = (a >= 0).float() * torch.exp(-a * math.log(2) / 0.1) a = a * ha.float().view(-1, 1).expand_as(a) / (1 + args.nb_classes) noise = a.new(a.size()).normal_(0, noise_level) a = a + noise - pos = torch.empty(nb, device = device).uniform_(0.0, 0.5) - b1 = torch.linspace(0, 1, length, device = device).view(1, -1).expand(nb, -1) + pos = torch.empty(nb, device=device).uniform_(0.0, 0.5) + b1 = torch.linspace(0, 1, length, device=device).view(1, -1).expand(nb, -1) b1 = b1 - pos.view(nb, 1) b1 = (b1 >= 0).float() * torch.exp(-b1 * math.log(2) / 0.1) * 0.25 pos = pos + hb.float() / (args.nb_classes + 1) * 0.5 # pos += pos.new(hb.size()).uniform_(0.0, 0.01) - b2 = torch.linspace(0, 1, length, device = device).view(1, -1).expand(nb, -1) + b2 = torch.linspace(0, 1, length, device=device).view(1, -1).expand(nb, -1) b2 = b2 - pos.view(nb, 1) b2 = (b2 >= 0).float() * torch.exp(-b2 * math.log(2) / 0.1) * 0.25 @@ -222,29 +237,33 @@ def create_sequences_pairs(train = False): return a, b, ha + ###################################################################### + class NetForImagePair(nn.Module): def __init__(self): super().__init__() self.features_a = nn.Sequential( - nn.Conv2d(1, 16, kernel_size = 5), - nn.MaxPool2d(3), nn.ReLU(), - nn.Conv2d(16, 32, kernel_size = 5), - nn.MaxPool2d(2), nn.ReLU(), + nn.Conv2d(1, 16, kernel_size=5), + nn.MaxPool2d(3), + nn.ReLU(), + nn.Conv2d(16, 32, kernel_size=5), + nn.MaxPool2d(2), + nn.ReLU(), ) self.features_b = nn.Sequential( - nn.Conv2d(1, 16, kernel_size = 5), - nn.MaxPool2d(3), nn.ReLU(), - nn.Conv2d(16, 32, kernel_size = 5), - nn.MaxPool2d(2), nn.ReLU(), + nn.Conv2d(1, 16, kernel_size=5), + nn.MaxPool2d(3), + nn.ReLU(), + nn.Conv2d(16, 32, kernel_size=5), + nn.MaxPool2d(2), + nn.ReLU(), ) self.fully_connected = nn.Sequential( - nn.Linear(256, 200), - nn.ReLU(), - nn.Linear(200, 1) + nn.Linear(256, 200), nn.ReLU(), nn.Linear(200, 1) ) def forward(self, a, b): @@ -253,28 +272,33 @@ class NetForImagePair(nn.Module): x = torch.cat((a, b), 1) return self.fully_connected(x) + ###################################################################### + class NetForImageValuesPair(nn.Module): def __init__(self): super().__init__() self.features_a = nn.Sequential( - nn.Conv2d(1, 16, kernel_size = 5), - nn.MaxPool2d(3), nn.ReLU(), - nn.Conv2d(16, 32, kernel_size = 5), - nn.MaxPool2d(2), nn.ReLU(), + nn.Conv2d(1, 16, kernel_size=5), + nn.MaxPool2d(3), + nn.ReLU(), + nn.Conv2d(16, 32, kernel_size=5), + nn.MaxPool2d(2), + nn.ReLU(), ) self.features_b = nn.Sequential( - nn.Linear(2, 32), nn.ReLU(), - nn.Linear(32, 32), nn.ReLU(), - nn.Linear(32, 128), nn.ReLU(), + nn.Linear(2, 32), + nn.ReLU(), + nn.Linear(32, 32), + nn.ReLU(), + nn.Linear(32, 128), + nn.ReLU(), ) self.fully_connected = nn.Sequential( - nn.Linear(256, 200), - nn.ReLU(), - nn.Linear(200, 1) + nn.Linear(256, 200), nn.ReLU(), nn.Linear(200, 1) ) def forward(self, a, b): @@ -283,24 +307,25 @@ class NetForImageValuesPair(nn.Module): x = torch.cat((a, b), 1) return self.fully_connected(x) + ###################################################################### -class NetForSequencePair(nn.Module): +class NetForSequencePair(nn.Module): def feature_model(self): kernel_size = 11 pooling_size = 4 - return nn.Sequential( - nn.Conv1d( 1, self.nc, kernel_size = kernel_size), + return nn.Sequential( + nn.Conv1d(1, self.nc, kernel_size=kernel_size), nn.AvgPool1d(pooling_size), nn.LeakyReLU(), - nn.Conv1d(self.nc, self.nc, kernel_size = kernel_size), + nn.Conv1d(self.nc, self.nc, kernel_size=kernel_size), nn.AvgPool1d(pooling_size), nn.LeakyReLU(), - nn.Conv1d(self.nc, self.nc, kernel_size = kernel_size), + nn.Conv1d(self.nc, self.nc, kernel_size=kernel_size), nn.AvgPool1d(pooling_size), nn.LeakyReLU(), - nn.Conv1d(self.nc, self.nc, kernel_size = kernel_size), + nn.Conv1d(self.nc, self.nc, kernel_size=kernel_size), nn.AvgPool1d(pooling_size), nn.LeakyReLU(), ) @@ -315,9 +340,7 @@ class NetForSequencePair(nn.Module): self.features_b = self.feature_model() self.fully_connected = nn.Sequential( - nn.Linear(2 * self.nc, self.nh), - nn.ReLU(), - nn.Linear(self.nh, 1) + nn.Linear(2 * self.nc, self.nh), nn.ReLU(), nn.Linear(self.nh, 1) ) def forward(self, a, b): @@ -332,17 +355,18 @@ class NetForSequencePair(nn.Module): x = torch.cat((a.view(a.size(0), -1), b.view(b.size(0), -1)), 1) return self.fully_connected(x) + ###################################################################### -if args.data == 'image_pair': +if args.data == "image_pair": create_pairs = create_image_pairs model = NetForImagePair() -elif args.data == 'image_values_pair': +elif args.data == "image_values_pair": create_pairs = create_image_values_pairs model = NetForImageValuesPair() -elif args.data == 'sequence_pair': +elif args.data == "sequence_pair": create_pairs = create_sequences_pairs model = NetForSequencePair() @@ -350,65 +374,70 @@ elif args.data == 'sequence_pair': ## Save for figures a, b, c = create_pairs() for k in range(10): - file = open(f'train_{k:02d}.dat', 'w') + file = open(f"train_{k:02d}.dat", "w") for i in range(a.size(1)): - file.write(f'{a[k, i]:f} {b[k,i]:f}\n') + file.write(f"{a[k, i]:f} {b[k,i]:f}\n") file.close() ###################### else: - raise Exception('Unknown data ' + args.data) + raise Exception("Unknown data " + args.data) ###################################################################### # Train -print(f'nb_parameters {sum(x.numel() for x in model.parameters())}') +print(f"nb_parameters {sum(x.numel() for x in model.parameters())}") model.to(device) -input_a, input_b, classes = create_pairs(train = True) +input_a, input_b, classes = create_pairs(train=True) for e in range(args.nb_epochs): - - optimizer = torch.optim.Adam(model.parameters(), lr = args.learning_rate) + optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate) input_br = input_b[torch.randperm(input_b.size(0))] acc_mi = 0.0 - for batch_a, batch_b, batch_br in zip(input_a.split(args.batch_size), - input_b.split(args.batch_size), - input_br.split(args.batch_size)): - mi = model(batch_a, batch_b).mean() - model(batch_a, batch_br).exp().mean().log() + for batch_a, batch_b, batch_br in zip( + input_a.split(args.batch_size), + input_b.split(args.batch_size), + input_br.split(args.batch_size), + ): + mi = ( + model(batch_a, batch_b).mean() - model(batch_a, batch_br).exp().mean().log() + ) acc_mi += mi.item() - loss = - mi + loss = -mi optimizer.zero_grad() loss.backward() optimizer.step() - acc_mi /= (input_a.size(0) // args.batch_size) + acc_mi /= input_a.size(0) // args.batch_size - print(f'{e+1} {acc_mi / math.log(2):.04f} {entropy(classes) / math.log(2):.04f}') + print(f"{e+1} {acc_mi / math.log(2):.04f} {entropy(classes) / math.log(2):.04f}") sys.stdout.flush() ###################################################################### # Test -input_a, input_b, classes = create_pairs(train = False) +input_a, input_b, classes = create_pairs(train=False) input_br = input_b[torch.randperm(input_b.size(0))] acc_mi = 0.0 -for batch_a, batch_b, batch_br in zip(input_a.split(args.batch_size), - input_b.split(args.batch_size), - input_br.split(args.batch_size)): +for batch_a, batch_b, batch_br in zip( + input_a.split(args.batch_size), + input_b.split(args.batch_size), + input_br.split(args.batch_size), +): mi = model(batch_a, batch_b).mean() - model(batch_a, batch_br).exp().mean().log() acc_mi += mi.item() -acc_mi /= (input_a.size(0) // args.batch_size) +acc_mi /= input_a.size(0) // args.batch_size -print(f'test {acc_mi / math.log(2):.04f} {entropy(classes) / math.log(2):.04f}') +print(f"test {acc_mi / math.log(2):.04f} {entropy(classes) / math.log(2):.04f}") ###################################################################### diff --git a/minidiffusion.py b/minidiffusion.py index 066cbbb..a3ffda0 100755 --- a/minidiffusion.py +++ b/minidiffusion.py @@ -13,22 +13,25 @@ import torch, torchvision from torch import nn from torch.nn import functional as F -device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -print(f'device {device}') +print(f"device {device}") ###################################################################### + def sample_gaussian_mixture(nb): p, std = 0.3, 0.2 result = torch.randn(nb, 1) * std result = result + torch.sign(torch.rand(result.size()) - p) / 2 return result + def sample_ramp(nb): result = torch.min(torch.rand(nb, 1), torch.rand(nb, 1)) return result + def sample_two_discs(nb): a = torch.rand(nb) * math.pi * 2 b = torch.rand(nb).sqrt() @@ -39,6 +42,7 @@ def sample_two_discs(nb): result[:, 1] = a.sin() * b - 0.5 + q return result + def sample_disc_grid(nb): a = torch.rand(nb) * math.pi * 2 b = torch.rand(nb).sqrt() @@ -51,6 +55,7 @@ def sample_disc_grid(nb): result[:, 1] = a.sin() * b + r return result + def sample_spiral(nb): u = torch.rand(nb) rho = u * 0.65 + 0.25 + torch.rand(nb) * 0.15 @@ -60,13 +65,16 @@ def sample_spiral(nb): result[:, 1] = theta.sin() * rho return result + def sample_mnist(nb): - train_set = torchvision.datasets.MNIST(root = './data/', train = True, download = True) + train_set = torchvision.datasets.MNIST(root="./data/", train=True, download=True) result = train_set.data[:nb].to(device).view(-1, 1, 28, 28).float() return result + samplers = { - f.__name__.removeprefix('sample_') : f for f in [ + f.__name__.removeprefix("sample_"): f + for f in [ sample_gaussian_mixture, sample_ramp, sample_two_discs, @@ -79,45 +87,40 @@ samplers = { ###################################################################### parser = argparse.ArgumentParser( - description = '''A minimal implementation of Jonathan Ho, Ajay Jain, Pieter Abbeel + description="""A minimal implementation of Jonathan Ho, Ajay Jain, Pieter Abbeel "Denoising Diffusion Probabilistic Models" (2020) -https://arxiv.org/abs/2006.11239''', - - formatter_class = argparse.ArgumentDefaultsHelpFormatter +https://arxiv.org/abs/2006.11239""", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, ) -parser.add_argument('--seed', - type = int, default = 0, - help = 'Random seed, < 0 is no seeding') +parser.add_argument( + "--seed", type=int, default=0, help="Random seed, < 0 is no seeding" +) -parser.add_argument('--nb_epochs', - type = int, default = 100, - help = 'How many epochs') +parser.add_argument("--nb_epochs", type=int, default=100, help="How many epochs") -parser.add_argument('--batch_size', - type = int, default = 25, - help = 'Batch size') +parser.add_argument("--batch_size", type=int, default=25, help="Batch size") -parser.add_argument('--nb_samples', - type = int, default = 25000, - help = 'Number of training examples') +parser.add_argument( + "--nb_samples", type=int, default=25000, help="Number of training examples" +) -parser.add_argument('--learning_rate', - type = float, default = 1e-3, - help = 'Learning rate') +parser.add_argument("--learning_rate", type=float, default=1e-3, help="Learning rate") -parser.add_argument('--ema_decay', - type = float, default = 0.9999, - help = 'EMA decay, <= 0 is no EMA') +parser.add_argument( + "--ema_decay", type=float, default=0.9999, help="EMA decay, <= 0 is no EMA" +) -data_list = ', '.join( [ str(k) for k in samplers ]) +data_list = ", ".join([str(k) for k in samplers]) -parser.add_argument('--data', - type = str, default = 'gaussian_mixture', - help = f'Toy data-set to use: {data_list}') +parser.add_argument( + "--data", + type=str, + default="gaussian_mixture", + help=f"Toy data-set to use: {data_list}", +) -parser.add_argument('--no_window', - action='store_true', default = False) +parser.add_argument("--no_window", action="store_true", default=False) args = parser.parse_args() @@ -131,11 +134,12 @@ if args.seed >= 0: ###################################################################### + class EMA: def __init__(self, model, decay): self.model = model self.decay = decay - self.mem = { } + self.mem = {} with torch.no_grad(): for p in model.parameters(): self.mem[p] = p.clone() @@ -150,11 +154,13 @@ class EMA: for p in self.model.parameters(): p.copy_(self.mem[p]) + ###################################################################### # Gets a pair (x, t) and appends t (scalar or 1d tensor) to x as an # additional dimension / channel + class TimeAppender(nn.Module): def __init__(self): super().__init__() @@ -163,9 +169,10 @@ class TimeAppender(nn.Module): x, t = u if not torch.is_tensor(t): t = x.new_full((x.size(0),), t) - t = t.view((-1,) + (1,) * (x.dim() - 1)).expand_as(x[:,:1]) + t = t.view((-1,) + (1,) * (x.dim() - 1)).expand_as(x[:, :1]) return torch.cat((x, t), 1) + class ConvNet(nn.Module): def __init__(self, in_channels, out_channels): super().__init__() @@ -174,29 +181,30 @@ class ConvNet(nn.Module): self.core = nn.Sequential( TimeAppender(), - nn.Conv2d(in_channels + 1, nc, ks, padding = ks//2), + nn.Conv2d(in_channels + 1, nc, ks, padding=ks // 2), nn.ReLU(), - nn.Conv2d(nc, nc, ks, padding = ks//2), + nn.Conv2d(nc, nc, ks, padding=ks // 2), nn.ReLU(), - nn.Conv2d(nc, nc, ks, padding = ks//2), + nn.Conv2d(nc, nc, ks, padding=ks // 2), nn.ReLU(), - nn.Conv2d(nc, nc, ks, padding = ks//2), + nn.Conv2d(nc, nc, ks, padding=ks // 2), nn.ReLU(), - nn.Conv2d(nc, nc, ks, padding = ks//2), + nn.Conv2d(nc, nc, ks, padding=ks // 2), nn.ReLU(), - nn.Conv2d(nc, out_channels, ks, padding = ks//2), + nn.Conv2d(nc, out_channels, ks, padding=ks // 2), ) def forward(self, u): return self.core(u) + ###################################################################### # Data try: train_input = samplers[args.data](args.nb_samples).to(device) except KeyError: - print(f'unknown data {args.data}') + print(f"unknown data {args.data}") exit(1) train_mean, train_std = train_input.mean(), train_input.std() @@ -219,52 +227,53 @@ if train_input.dim() == 2: ) elif train_input.dim() == 4: - model = ConvNet(train_input.size(1), train_input.size(1)) model.to(device) -print(f'nb_parameters {sum([ p.numel() for p in model.parameters() ])}') +print(f"nb_parameters {sum([ p.numel() for p in model.parameters() ])}") ###################################################################### # Generate -def generate(size, T, alpha, alpha_bar, sigma, model, train_mean, train_std): +def generate(size, T, alpha, alpha_bar, sigma, model, train_mean, train_std): with torch.no_grad(): + x = torch.randn(size, device=device) - x = torch.randn(size, device = device) - - for t in range(T-1, -1, -1): + for t in range(T - 1, -1, -1): output = model((x, t / (T - 1) - 0.5)) z = torch.zeros_like(x) if t == 0 else torch.randn_like(x) - x = 1/torch.sqrt(alpha[t]) \ - * (x - (1-alpha[t]) / torch.sqrt(1-alpha_bar[t]) * output) \ + x = ( + 1 + / torch.sqrt(alpha[t]) + * (x - (1 - alpha[t]) / torch.sqrt(1 - alpha_bar[t]) * output) + sigma[t] * z + ) x = x * train_std + train_mean return x + ###################################################################### # Train T = 1000 -beta = torch.linspace(1e-4, 0.02, T, device = device) +beta = torch.linspace(1e-4, 0.02, T, device=device) alpha = 1 - beta alpha_bar = alpha.log().cumsum(0).exp() sigma = beta.sqrt() -ema = EMA(model, decay = args.ema_decay) if args.ema_decay > 0 else None +ema = EMA(model, decay=args.ema_decay) if args.ema_decay > 0 else None for k in range(args.nb_epochs): - acc_loss = 0 - optimizer = torch.optim.Adam(model.parameters(), lr = args.learning_rate) + optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate) for x0 in train_input.split(args.batch_size): x0 = (x0 - train_mean) / train_std - t = torch.randint(T, (x0.size(0),) + (1,) * (x0.dim() - 1), device = x0.device) + t = torch.randint(T, (x0.size(0),) + (1,) * (x0.dim() - 1), device=x0.device) eps = torch.randn_like(x0) xt = torch.sqrt(alpha_bar[t]) * x0 + torch.sqrt(1 - alpha_bar[t]) * eps output = model((xt, t / (T - 1) - 0.5)) @@ -275,11 +284,13 @@ for k in range(args.nb_epochs): loss.backward() optimizer.step() - if ema is not None: ema.step() + if ema is not None: + ema.step() - print(f'{k} {acc_loss / train_input.size(0)}') + print(f"{k} {acc_loss / train_input.size(0)}") -if ema is not None: ema.copy_to_model() +if ema is not None: + ema.copy_to_model() ###################################################################### # Plot @@ -289,100 +300,107 @@ model.eval() ######################################## # Nx1 -> histogram if train_input.dim() == 2 and train_input.size(1) == 1: - fig = plt.figure() fig.set_figheight(5) fig.set_figwidth(8) ax = fig.add_subplot(1, 1, 1) - x = generate((10000, 1), T, alpha, alpha_bar, sigma, - model, train_mean, train_std) + x = generate((10000, 1), T, alpha, alpha_bar, sigma, model, train_mean, train_std) ax.set_xlim(-1.25, 1.25) ax.spines.right.set_visible(False) ax.spines.top.set_visible(False) - d = train_input.flatten().detach().to('cpu').numpy() - ax.hist(d, 25, (-1, 1), - density = True, - histtype = 'bar', edgecolor = 'white', color = 'lightblue', label = 'Train') + d = train_input.flatten().detach().to("cpu").numpy() + ax.hist( + d, + 25, + (-1, 1), + density=True, + histtype="bar", + edgecolor="white", + color="lightblue", + label="Train", + ) - d = x.flatten().detach().to('cpu').numpy() - ax.hist(d, 25, (-1, 1), - density = True, - histtype = 'step', color = 'red', label = 'Synthesis') + d = x.flatten().detach().to("cpu").numpy() + ax.hist( + d, 25, (-1, 1), density=True, histtype="step", color="red", label="Synthesis" + ) - ax.legend(frameon = False, loc = 2) + ax.legend(frameon=False, loc=2) - filename = f'minidiffusion_{args.data}.pdf' - print(f'saving {filename}') - fig.savefig(filename, bbox_inches='tight') + filename = f"minidiffusion_{args.data}.pdf" + print(f"saving {filename}") + fig.savefig(filename, bbox_inches="tight") - if not args.no_window and hasattr(plt.get_current_fig_manager(), 'window'): + if not args.no_window and hasattr(plt.get_current_fig_manager(), "window"): plt.get_current_fig_manager().window.setGeometry(2, 2, 1024, 768) plt.show() ######################################## # Nx2 -> scatter plot elif train_input.dim() == 2 and train_input.size(1) == 2: - fig = plt.figure() fig.set_figheight(6) fig.set_figwidth(6) ax = fig.add_subplot(1, 1, 1) - x = generate((1000, 2), T, alpha, alpha_bar, sigma, - model, train_mean, train_std) + x = generate((1000, 2), T, alpha, alpha_bar, sigma, model, train_mean, train_std) ax.set_xlim(-1.5, 1.5) ax.set_ylim(-1.5, 1.5) - ax.set(aspect = 1) + ax.set(aspect=1) ax.spines.right.set_visible(False) ax.spines.top.set_visible(False) - d = train_input[:x.size(0)].detach().to('cpu').numpy() - ax.scatter(d[:, 0], d[:, 1], - s = 2.5, color = 'gray', label = 'Train') + d = train_input[: x.size(0)].detach().to("cpu").numpy() + ax.scatter(d[:, 0], d[:, 1], s=2.5, color="gray", label="Train") - d = x.detach().to('cpu').numpy() - ax.scatter(d[:, 0], d[:, 1], - s = 2.0, color = 'red', label = 'Synthesis') + d = x.detach().to("cpu").numpy() + ax.scatter(d[:, 0], d[:, 1], s=2.0, color="red", label="Synthesis") - ax.legend(frameon = False, loc = 2) + ax.legend(frameon=False, loc=2) - filename = f'minidiffusion_{args.data}.pdf' - print(f'saving {filename}') - fig.savefig(filename, bbox_inches='tight') + filename = f"minidiffusion_{args.data}.pdf" + print(f"saving {filename}") + fig.savefig(filename, bbox_inches="tight") - if not args.no_window and hasattr(plt.get_current_fig_manager(), 'window'): + if not args.no_window and hasattr(plt.get_current_fig_manager(), "window"): plt.get_current_fig_manager().window.setGeometry(2, 2, 1024, 768) plt.show() ######################################## # NxCxHxW -> image elif train_input.dim() == 4: + x = generate( + (128,) + train_input.size()[1:], + T, + alpha, + alpha_bar, + sigma, + model, + train_mean, + train_std, + ) - x = generate((128,) + train_input.size()[1:], T, alpha, alpha_bar, sigma, - model, train_mean, train_std) - - x = torchvision.utils.make_grid(x.clamp(min = 0, max = 255), - nrow = 16, padding = 1, pad_value = 64) - x = F.pad(x, pad = (2, 2, 2, 2), value = 64)[None] + x = torchvision.utils.make_grid( + x.clamp(min=0, max=255), nrow=16, padding=1, pad_value=64 + ) + x = F.pad(x, pad=(2, 2, 2, 2), value=64)[None] - t = torchvision.utils.make_grid(train_input[:128], - nrow = 16, padding = 1, pad_value = 64) - t = F.pad(t, pad = (2, 2, 2, 2), value = 64)[None] + t = torchvision.utils.make_grid(train_input[:128], nrow=16, padding=1, pad_value=64) + t = F.pad(t, pad=(2, 2, 2, 2), value=64)[None] result = 1 - torch.cat((t, x), 2) / 255 - filename = f'minidiffusion_{args.data}.png' - print(f'saving {filename}') + filename = f"minidiffusion_{args.data}.png" + print(f"saving {filename}") torchvision.utils.save_image(result, filename) else: - - print(f'cannot plot result of size {train_input.size()}') + print(f"cannot plot result of size {train_input.size()}") ###################################################################### diff --git a/miniflow.py b/miniflow.py index 04b9a23..ad7b431 100755 --- a/miniflow.py +++ b/miniflow.py @@ -19,33 +19,40 @@ from torch.nn import functional as F ###################################################################### + def phi(x): p, std = 0.3, 0.2 - mu = (1 - p) * torch.exp(LogProba((x - 0.5) / std, math.log(1 / std))) + \ - p * torch.exp(LogProba((x + 0.5) / std, math.log(1 / std))) + mu = (1 - p) * torch.exp( + LogProba((x - 0.5) / std, math.log(1 / std)) + ) + p * torch.exp(LogProba((x + 0.5) / std, math.log(1 / std))) return mu + def sample_phi(nb): p, std = 0.3, 0.2 result = torch.empty(nb).normal_(0, std) result = result + torch.sign(torch.rand(result.size()) - p) / 2 return result + ###################################################################### + def LogProba(x, ldj): - log_p = ldj - 0.5 * (x**2 + math.log(2*pi)) + log_p = ldj - 0.5 * (x**2 + math.log(2 * pi)) return log_p + ###################################################################### + class PiecewiseLinear(nn.Module): def __init__(self, nb, xmin, xmax): super().__init__() self.xmin = xmin self.xmax = xmax self.nb = nb - self.alpha = nn.Parameter(torch.tensor([xmin], dtype = torch.float)) + self.alpha = nn.Parameter(torch.tensor([xmin], dtype=torch.float)) mu = math.log((xmax - xmin) / nb) self.xi = nn.Parameter(torch.empty(nb + 1).normal_(mu, 1e-4)) @@ -64,9 +71,12 @@ class PiecewiseLinear(nn.Module): assert (y >= ys[0, 0]).min() and (y <= ys[0, self.nb]).min() yk = ys[:, :-1] ykp1 = ys[:, 1:] - x = self.xmin + (self.xmax - self.xmin)/self.nb * ((yy >= yk) * (yy < ykp1).long() * (k + (yy - yk)/(ykp1 - yk))).sum(1) + x = self.xmin + (self.xmax - self.xmin) / self.nb * ( + (yy >= yk) * (yy < ykp1).long() * (k + (yy - yk) / (ykp1 - yk)) + ).sum(1) return x + ###################################################################### # Training @@ -74,11 +84,11 @@ nb_samples = 25000 nb_epochs = 250 batch_size = 100 -model = PiecewiseLinear(nb = 1001, xmin = -4, xmax = 4) +model = PiecewiseLinear(nb=1001, xmin=-4, xmax=4) train_input = sample_phi(nb_samples) -optimizer = torch.optim.Adam(model.parameters(), lr = 1e-4) +optimizer = torch.optim.Adam(model.parameters(), lr=1e-4) criterion = nn.MSELoss() for k in range(nb_epochs): @@ -88,19 +98,19 @@ for k in range(nb_epochs): input.requires_grad_() output = model(input) - derivatives, = autograd.grad( - output.sum(), input, - retain_graph = True, create_graph = True + (derivatives,) = autograd.grad( + output.sum(), input, retain_graph=True, create_graph=True ) - loss = ( 0.5 * (output**2 + math.log(2*pi)) - derivatives.log() ).mean() + loss = (0.5 * (output**2 + math.log(2 * pi)) - derivatives.log()).mean() optimizer.zero_grad() loss.backward() optimizer.step() acc_loss += loss.item() - if k%10 == 0: print(k, loss.item()) + if k % 10 == 0: + print(k, loss.item()) ###################################################################### @@ -131,33 +141,35 @@ ax = fig.add_subplot(1, 1, 1) # ax.set_ylim(-0.25, 1.25) # ax.axis('off') -ax.plot(input, output, '-', color = 'tab:red') +ax.plot(input, output, "-", color="tab:red") -filename = 'miniflow_mapping.pdf' -print(f'Saving {filename}') -fig.savefig(filename, bbox_inches='tight') +filename = "miniflow_mapping.pdf" +print(f"Saving {filename}") +fig.savefig(filename, bbox_inches="tight") # plt.show() ###################################################################### -green_dist = '#bfdfbf' +green_dist = "#bfdfbf" fig = plt.figure() ax = fig.add_subplot(1, 1, 1) # ax.set_xlim(-4.5, 4.5) # ax.set_ylim(-0.1, 1.1) -lines = list(([(x_in.item(), 0), (x_out.item(), 0.5)]) for (x_in, x_out) in zip(input, output)) -lc = mc.LineCollection(lines, color = 'tab:red', linewidth = 0.1) +lines = list( + ([(x_in.item(), 0), (x_out.item(), 0.5)]) for (x_in, x_out) in zip(input, output) +) +lc = mc.LineCollection(lines, color="tab:red", linewidth=0.1) ax.add_collection(lc) -ax.axis('off') +ax.axis("off") -ax.fill_between(input, 0.52, mu_N * 0.2 + 0.52, color = green_dist) -ax.fill_between(input, -0.30, mu * 0.2 - 0.30, color = green_dist) +ax.fill_between(input, 0.52, mu_N * 0.2 + 0.52, color=green_dist) +ax.fill_between(input, -0.30, mu * 0.2 - 0.30, color=green_dist) -filename = 'miniflow_flow.pdf' -print(f'Saving {filename}') -fig.savefig(filename, bbox_inches='tight') +filename = "miniflow_flow.pdf" +print(f"Saving {filename}") +fig.savefig(filename, bbox_inches="tight") # plt.show() @@ -165,16 +177,16 @@ fig.savefig(filename, bbox_inches='tight') fig = plt.figure() ax = fig.add_subplot(1, 1, 1) -ax.axis('off') +ax.axis("off") -ax.fill_between(input, 0, mu, color = green_dist) +ax.fill_between(input, 0, mu, color=green_dist) # ax.plot(input, mu, '-', color = 'tab:blue') # ax.step(input, mu_hat, '-', where='mid', color = 'tab:red') -ax.plot(input, mu_hat, '-', color = 'tab:red') +ax.plot(input, mu_hat, "-", color="tab:red") -filename = 'miniflow_dist.pdf' -print(f'Saving {filename}') -fig.savefig(filename, bbox_inches='tight') +filename = "miniflow_dist.pdf" +print(f"Saving {filename}") +fig.savefig(filename, bbox_inches="tight") # plt.show() @@ -182,15 +194,15 @@ fig.savefig(filename, bbox_inches='tight') fig = plt.figure() ax = fig.add_subplot(1, 1, 1) -ax.axis('off') +ax.axis("off") # ax.plot(input, mu, '-', color = 'tab:blue') -ax.fill_between(input, 0, mu, color = green_dist) +ax.fill_between(input, 0, mu, color=green_dist) # ax.step(input, mu_hat, '-', where='mid', color = 'tab:red') -filename = 'miniflow_target_dist.pdf' -print(f'Saving {filename}') -fig.savefig(filename, bbox_inches='tight') +filename = "miniflow_target_dist.pdf" +print(f"Saving {filename}") +fig.savefig(filename, bbox_inches="tight") # plt.show() diff --git a/poly.py b/poly.py index 1b157a8..8aac207 100755 --- a/poly.py +++ b/poly.py @@ -9,6 +9,7 @@ import torch + def pol_prod(a, b): m = a[:, None] * b[None, :] mm = m.new() @@ -16,23 +17,26 @@ def pol_prod(a, b): k = torch.arange(a.size(0))[:, None] + torch.arange(b.size(0))[None, :] kk = k.new() kk.set_(k.storage(), 0, (k.size(0), k.size(0) + k.size(1) - 1), (k.size(1) - 1, 1)) - q = (kk == torch.arange(a.size(0) + b.size(0) - 1)[None, :]) + q = kk == torch.arange(a.size(0) + b.size(0) - 1)[None, :] return (mm * q).sum(0) + def pol_eval(a, x): d = torch.arange(a.size(0)) return (x[:, None].pow(d[None, :]) * a[None, :]).sum(1) + def pol_prim(a): n = torch.arange(a.size(0) + 1).float() n[1:] = a / n[1:] return n + ###################################################################### -if __name__ == '__main__': - a = torch.tensor([1., 2., 3.]) - b = torch.tensor([2., 5.]) +if __name__ == "__main__": + a = torch.tensor([1.0, 2.0, 3.0]) + b = torch.tensor([2.0, 5.0]) print(pol_prod(a, b)) print(pol_prim(b)) print(pol_eval(a, torch.tensor([0.0, 1.0, 2.0]))) diff --git a/rmax.py b/rmax.py new file mode 100755 index 0000000..291ce92 --- /dev/null +++ b/rmax.py @@ -0,0 +1,40 @@ +#!/usr/bin/env python + +import torch + +################################################## + + +def rmax(x): + a = x.max(-1, keepdim=True) + i = torch.arange(x.size(-1) - 1)[None, :] + y = torch.cat( + ( + (i < a.indices) * (x - a.values)[:, :-1] + + (i >= a.indices) * (a.values - x)[:, 1:], + a.values, + ), + -1, + ) + return y + + +def rmax_back(y): + u = torch.nn.functional.pad(y, (1, -1)) + x = ( + (y < 0) * (y[:, -1:] + y) + + (y >= 0) * (u < 0) * (y[:, -1:]) + + (y >= 0) * (u >= 0) * (y[:, -1:] - u) + ) + return x + + +################################################## + +x = torch.randn(3, 14) +y = rmax(x) +print(f"{x.size()=} {x.max(-1).values=}") +print(f"{y.size()=} {y[:,-1]=}") + +z = rmax_back(y) +print(f"{(z-x).abs().max()=}") diff --git a/sizer.py b/sizer.py index dff36eb..5887e4a 100755 --- a/sizer.py +++ b/sizer.py @@ -13,7 +13,9 @@ from torch import nn ###################################################################### if len(sys.argv) < 2: - print(sys.argv[0] + ''' + print( + sys.argv[0] + + """ For example: @@ -24,7 +26,8 @@ nn.Conv2d(32, 32, 3, padding = 1) nn.MaxPool2d(2) nn.Conv2d(32, 64, 3, padding = 1) nn.MaxPool2d(5) -nn.Conv2d(64, 64, (3, 4))''') +nn.Conv2d(64, 64, (3, 4))""" + ) exit(1) ###################################################################### @@ -36,15 +39,15 @@ while True: t = os.stat(sys.argv[1])[stat.ST_MTIME] if t > pt: pt = t - os.system('clear') + os.system("clear") try: - temp = [l.strip('\n\r') for l in open(sys.argv[1], 'r').readlines()] + temp = [l.strip("\n\r") for l in open(sys.argv[1], "r").readlines()] x = torch.zeros(eval(temp.pop(0))) - print('-> ' + str(tuple(x.size()))) + print("-> " + str(tuple(x.size()))) for k in temp: - print(' ' + k) - x = eval(k + '(x)') - print('-> ' + str(tuple(x.size()))) + print(" " + k) + x = eval(k + "(x)") + print("-> " + str(tuple(x.size()))) except: - print('** Error **') + print("** Error **") time.sleep(1) diff --git a/speed.py b/speed.py index 075b07e..8363a6c 100755 --- a/speed.py +++ b/speed.py @@ -8,19 +8,19 @@ import time, torch if torch.cuda.is_available(): - device = torch.device('cuda') + device = torch.device("cuda") sync = torch.cuda.synchronize else: - device = torch.device('cpu') + device = torch.device("cpu") sync = lambda: None max_duration = 30 d1, d2, d3 = 2048, 2048, 2048 -for t in [ torch.float32, torch.float16 ]: +for t in [torch.float32, torch.float16]: try: - a = torch.rand(d1, d2, device = device, dtype = t) - b = torch.rand(d2, d3, device = device, dtype = t) + a = torch.rand(d1, d2, device=device, dtype=t) + b = torch.rand(d2, d3, device=device, dtype=t) nb_runs = 0 sync() @@ -31,15 +31,15 @@ for t in [ torch.float32, torch.float16 ]: sync() duration = time.perf_counter() - start_time - nb_flop = float(nb_runs * d1 * d2 * d3 * 2) # 1 multiply-and-add is 2 ops + nb_flop = float(nb_runs * d1 * d2 * d3 * 2) # 1 multiply-and-add is 2 ops speed = nb_flop / duration - for u in [ '', 'K', 'M', 'G', 'T', 'P' ]: - if speed < 1e3: break + for u in ["", "K", "M", "G", "T", "P"]: + if speed < 1e3: + break speed /= 1e3 - print(f'{speed:.02f} {u}flops with {t} on {device}') + print(f"{speed:.02f} {u}flops with {t} on {device}") except: - - print(f'{t} is not available on {device}') + print(f"{t} is not available on {device}") diff --git a/tensorstack.py b/tensorstack.py index c9a6c2f..42a051e 100755 --- a/tensorstack.py +++ b/tensorstack.py @@ -9,52 +9,53 @@ from torch import Tensor import sys + def exception_hook(exc_type, exc_value, tb): - r'''Hacks the call stack message to show all the local variables in - case of RuntimeError or ValueError, and prints tensors as shape, - dtype and device. + r"""Hacks the call stack message to show all the local variables + in case of RuntimeError, ValueError, or TypeError and prints + tensors as shape, dtype and device. - ''' + """ - repr_orig=Tensor.__repr__ - Tensor.__repr__=lambda x: f'{x.size()}:{x.dtype}:{x.device}' + repr_orig = Tensor.__repr__ + Tensor.__repr__ = lambda x: f"{x.size()}:{x.dtype}:{x.device}" while tb: - print('--------------------------------------------------\n') + print("--------------------------------------------------\n") filename = tb.tb_frame.f_code.co_filename name = tb.tb_frame.f_code.co_name line_no = tb.tb_lineno print(f' File "{filename}", line {line_no}, in {name}') - print(open(filename, 'r').readlines()[line_no-1]) + print(open(filename, "r").readlines()[line_no - 1]) - if exc_type in { RuntimeError, ValueError }: - for n,v in tb.tb_frame.f_locals.items(): - print(f' {n} -> {v}') + if exc_type in {RuntimeError, ValueError, TypeError}: + for n, v in tb.tb_frame.f_locals.items(): + print(f" {n} -> {v}") print() tb = tb.tb_next - Tensor.__repr__=repr_orig + Tensor.__repr__ = repr_orig + + print(f"{exc_type.__name__}: {exc_value}") - print(f'{exc_type.__name__}: {exc_value}') sys.excepthook = exception_hook ###################################################################### -if __name__ == '__main__': - +if __name__ == "__main__": import torch - def dummy(a,b): - print(a@b) + def dummy(a, b): + print(a @ b) - def blah(a,b): - c=b+b - dummy(a,c) + def blah(a, b): + c = b + b + dummy(a, c) - mmm=torch.randn(2,3) - xxx=torch.randn(3) - #print(xxx@mmm) - blah(mmm,xxx) - blah(xxx,mmm) + mmm = torch.randn(2, 3) + xxx = torch.randn(3) + # print(xxx@mmm) + blah(mmm, xxx) + blah(xxx, mmm) diff --git a/tinyae.py b/tinyae.py index 70484f1..1608786 100755 --- a/tinyae.py +++ b/tinyae.py @@ -14,74 +14,75 @@ from torch.nn import functional as F ###################################################################### -device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") ###################################################################### -parser = argparse.ArgumentParser(description = 'Tiny LeNet-like auto-encoder.') +parser = argparse.ArgumentParser(description="Tiny LeNet-like auto-encoder.") -parser.add_argument('--nb_epochs', - type = int, default = 25) +parser.add_argument("--nb_epochs", type=int, default=25) -parser.add_argument('--batch_size', - type = int, default = 100) +parser.add_argument("--batch_size", type=int, default=100) -parser.add_argument('--data_dir', - type = str, default = './data/') +parser.add_argument("--data_dir", type=str, default="./data/") -parser.add_argument('--log_filename', - type = str, default = 'train.log') +parser.add_argument("--log_filename", type=str, default="train.log") -parser.add_argument('--embedding_dim', - type = int, default = 8) +parser.add_argument("--embedding_dim", type=int, default=8) -parser.add_argument('--nb_channels', - type = int, default = 32) +parser.add_argument("--nb_channels", type=int, default=32) args = parser.parse_args() -log_file = open(args.log_filename, 'w') +log_file = open(args.log_filename, "w") ###################################################################### + def log_string(s): t = time.strftime("%Y-%m-%d_%H:%M:%S - ", time.localtime()) if log_file is not None: - log_file.write(t + s + '\n') + log_file.write(t + s + "\n") log_file.flush() print(t + s) sys.stdout.flush() + ###################################################################### + class AutoEncoder(nn.Module): def __init__(self, nb_channels, embedding_dim): super(AutoEncoder, self).__init__() self.encoder = nn.Sequential( - nn.Conv2d(1, nb_channels, kernel_size = 5), # to 24x24 - nn.ReLU(inplace = True), - nn.Conv2d(nb_channels, nb_channels, kernel_size = 5), # to 20x20 - nn.ReLU(inplace = True), - nn.Conv2d(nb_channels, nb_channels, kernel_size = 4, stride = 2), # to 9x9 - nn.ReLU(inplace = True), - nn.Conv2d(nb_channels, nb_channels, kernel_size = 3, stride = 2), # to 4x4 - nn.ReLU(inplace = True), - nn.Conv2d(nb_channels, embedding_dim, kernel_size = 4) + nn.Conv2d(1, nb_channels, kernel_size=5), # to 24x24 + nn.ReLU(inplace=True), + nn.Conv2d(nb_channels, nb_channels, kernel_size=5), # to 20x20 + nn.ReLU(inplace=True), + nn.Conv2d(nb_channels, nb_channels, kernel_size=4, stride=2), # to 9x9 + nn.ReLU(inplace=True), + nn.Conv2d(nb_channels, nb_channels, kernel_size=3, stride=2), # to 4x4 + nn.ReLU(inplace=True), + nn.Conv2d(nb_channels, embedding_dim, kernel_size=4), ) self.decoder = nn.Sequential( - nn.ConvTranspose2d(embedding_dim, nb_channels, kernel_size = 4), - nn.ReLU(inplace = True), - nn.ConvTranspose2d(nb_channels, nb_channels, kernel_size = 3, stride = 2), # from 4x4 - nn.ReLU(inplace = True), - nn.ConvTranspose2d(nb_channels, nb_channels, kernel_size = 4, stride = 2), # from 9x9 - nn.ReLU(inplace = True), - nn.ConvTranspose2d(nb_channels, nb_channels, kernel_size = 5), # from 20x20 - nn.ReLU(inplace = True), - nn.ConvTranspose2d(nb_channels, 1, kernel_size = 5), # from 24x24 + nn.ConvTranspose2d(embedding_dim, nb_channels, kernel_size=4), + nn.ReLU(inplace=True), + nn.ConvTranspose2d( + nb_channels, nb_channels, kernel_size=3, stride=2 + ), # from 4x4 + nn.ReLU(inplace=True), + nn.ConvTranspose2d( + nb_channels, nb_channels, kernel_size=4, stride=2 + ), # from 9x9 + nn.ReLU(inplace=True), + nn.ConvTranspose2d(nb_channels, nb_channels, kernel_size=5), # from 20x20 + nn.ReLU(inplace=True), + nn.ConvTranspose2d(nb_channels, 1, kernel_size=5), # from 24x24 ) def encode(self, x): @@ -95,20 +96,23 @@ class AutoEncoder(nn.Module): x = self.decoder(x) return x + ###################################################################### -train_set = torchvision.datasets.MNIST(args.data_dir + '/mnist/', - train = True, download = True) +train_set = torchvision.datasets.MNIST( + args.data_dir + "/mnist/", train=True, download=True +) train_input = train_set.data.view(-1, 1, 28, 28).float() -test_set = torchvision.datasets.MNIST(args.data_dir + '/mnist/', - train = False, download = True) +test_set = torchvision.datasets.MNIST( + args.data_dir + "/mnist/", train=False, download=True +) test_input = test_set.data.view(-1, 1, 28, 28).float() ###################################################################### model = AutoEncoder(args.nb_channels, args.embedding_dim) -optimizer = optim.Adam(model.parameters(), lr = 1e-3) +optimizer = optim.Adam(model.parameters(), lr=1e-3) model.to(device) @@ -121,7 +125,6 @@ test_input.sub_(mu).div_(std) ###################################################################### for epoch in range(args.nb_epochs): - acc_loss = 0 for input in train_input.split(args.batch_size): @@ -134,7 +137,7 @@ for epoch in range(args.nb_epochs): acc_loss += loss.item() - log_string('acc_loss {:d} {:f}.'.format(epoch, acc_loss)) + log_string("acc_loss {:d} {:f}.".format(epoch, acc_loss)) ###################################################################### @@ -145,8 +148,8 @@ input = test_input[:256] z = model.encode(input) output = model.decode(z) -torchvision.utils.save_image(1 - input, 'ae-input.png', nrow = 16, pad_value = 0.8) -torchvision.utils.save_image(1 - output, 'ae-output.png', nrow = 16, pad_value = 0.8) +torchvision.utils.save_image(1 - input, "ae-input.png", nrow=16, pad_value=0.8) +torchvision.utils.save_image(1 - output, "ae-output.png", nrow=16, pad_value=0.8) # Dumb synthesis @@ -155,6 +158,6 @@ mu, std = z.mean(0), z.std(0) z = z.normal_() * std + mu output = model.decode(z) -torchvision.utils.save_image(1 - output, 'ae-synth.png', nrow = 16, pad_value = 0.8) +torchvision.utils.save_image(1 - output, "ae-synth.png", nrow=16, pad_value=0.8) ###################################################################### diff --git a/tinymnist.py b/tinymnist.py index 8642b22..896477e 100755 --- a/tinymnist.py +++ b/tinymnist.py @@ -12,47 +12,49 @@ from torch.nn import functional as F lr, nb_epochs, batch_size = 1e-1, 10, 100 -data_dir = os.environ.get('PYTORCH_DATA_DIR') or './data/' +data_dir = os.environ.get("PYTORCH_DATA_DIR") or "./data/" -device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") ###################################################################### -train_set = torchvision.datasets.MNIST(root = data_dir, train = True, download = True) +train_set = torchvision.datasets.MNIST(root=data_dir, train=True, download=True) train_input = train_set.data.view(-1, 1, 28, 28).float() train_targets = train_set.targets -test_set = torchvision.datasets.MNIST(root = data_dir, train = False, download = True) +test_set = torchvision.datasets.MNIST(root=data_dir, train=False, download=True) test_input = test_set.data.view(-1, 1, 28, 28).float() test_targets = test_set.targets ###################################################################### + class SomeLeNet(nn.Module): def __init__(self): super().__init__() - self.conv1 = nn.Conv2d(1, 32, kernel_size = 5) - self.conv2 = nn.Conv2d(32, 64, kernel_size = 5) + self.conv1 = nn.Conv2d(1, 32, kernel_size=5) + self.conv2 = nn.Conv2d(32, 64, kernel_size=5) self.fc1 = nn.Linear(256, 200) self.fc2 = nn.Linear(200, 10) def forward(self, x): - x = F.relu(F.max_pool2d(self.conv1(x), kernel_size = 3)) - x = F.relu(F.max_pool2d(self.conv2(x), kernel_size = 2)) + x = F.relu(F.max_pool2d(self.conv1(x), kernel_size=3)) + x = F.relu(F.max_pool2d(self.conv2(x), kernel_size=2)) x = x.view(x.size(0), -1) x = F.relu(self.fc1(x)) x = self.fc2(x) return x + ###################################################################### model = SomeLeNet() nb_parameters = sum(p.numel() for p in model.parameters()) -print(f'nb_parameters {nb_parameters}') +print(f"nb_parameters {nb_parameters}") -optimizer = torch.optim.SGD(model.parameters(), lr = lr) +optimizer = torch.optim.SGD(model.parameters(), lr=lr) criterion = nn.CrossEntropyLoss() model.to(device) @@ -68,10 +70,11 @@ test_input.sub_(mu).div_(std) start_time = time.perf_counter() for k in range(nb_epochs): - acc_loss = 0. + acc_loss = 0.0 - for input, targets in zip(train_input.split(batch_size), - train_targets.split(batch_size)): + for input, targets in zip( + train_input.split(batch_size), train_targets.split(batch_size) + ): output = model(input) loss = criterion(output, targets) acc_loss += loss.item() @@ -81,13 +84,14 @@ for k in range(nb_epochs): optimizer.step() nb_test_errors = 0 - for input, targets in zip(test_input.split(batch_size), - test_targets.split(batch_size)): + for input, targets in zip( + test_input.split(batch_size), test_targets.split(batch_size) + ): wta = model(input).argmax(1) nb_test_errors += (wta != targets).long().sum() test_error = nb_test_errors / test_input.size(0) duration = time.perf_counter() - start_time - print(f'loss {k} {duration:.02f}s {acc_loss:.02f} {test_error*100:.02f}%') + print(f"loss {k} {duration:.02f}s {acc_loss:.02f} {test_error*100:.02f}%") ###################################################################### -- 2.39.5