######################################################################
+
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:
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 = []
######################################################################
-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()
######################################################################
-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")
######################################################################
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)
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)
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)
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)
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)
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)
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)
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)
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:
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")
######################################################################
######################################################################
-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",
+)
######################################################################
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))))
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))))
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))
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)
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
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:
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()
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()
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)
######################################################################
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()
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()
######################################################################
+
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())
######################################################################
-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()
######################################################################
+
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)
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)
######################################################################
+
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)
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):
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
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)
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]
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, :]))
######################################################################
-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
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):
######################################################################
-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)
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",
)
)
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)
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():
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)
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)
######################################################################
-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
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()
for n, x in model.named_parameters():
print(n, x.size())
-
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")
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:
(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",
+)
######################################################################
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)
# 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:
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])
return a, b, c
+
######################################################################
# Returns a triplet a, b, c where a are the standard MNIST images, c
# 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:
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()
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
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):
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):
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(),
)
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):
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()
## 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}")
######################################################################
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()
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()
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
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,
######################################################################
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()
######################################################################
+
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()
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__()
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__()
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()
)
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))
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
########################################
# 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()}")
######################################################################
######################################################################
+
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))
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
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):
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())
######################################################################
# 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()
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()
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()
import torch
+
def pol_prod(a, b):
m = a[:, None] * b[None, :]
mm = m.new()
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])))
--- /dev/null
+#!/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()=}")
######################################################################
if len(sys.argv) < 2:
- print(sys.argv[0] + ''' <file to monitor>
+ print(
+ sys.argv[0]
+ + """ <file to monitor>
For example:
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)
######################################################################
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)
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()
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}")
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)
######################################################################
-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):
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)
######################################################################
for epoch in range(args.nb_epochs):
-
acc_loss = 0
for input in train_input.split(args.batch_size):
acc_loss += loss.item()
- log_string('acc_loss {:d} {:f}.'.format(epoch, acc_loss))
+ log_string("acc_loss {:d} {:f}.".format(epoch, acc_loss))
######################################################################
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
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)
######################################################################
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)
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()
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}%")
######################################################################
projects / pytorch.git / commitdiff