projects
/
pytorch.git
/ commitdiff
commit
grep
author
committer
pickaxe
?
search:
re
summary
|
shortlog
|
log
|
commit
| commitdiff |
tree
raw
|
patch
|
inline
| side by side (from parent 1:
9898347
)
Update.
author
Francois Fleuret
<francois@fleuret.org>
Tue, 16 Aug 2022 20:51:29 +0000
(22:51 +0200)
committer
Francois Fleuret
<francois@fleuret.org>
Tue, 16 Aug 2022 20:51:29 +0000
(22:51 +0200)
minidiffusion.py
patch
|
blob
|
history
diff --git
a/minidiffusion.py
b/minidiffusion.py
index
27842d9
..
7327522
100755
(executable)
--- a/
minidiffusion.py
+++ b/
minidiffusion.py
@@
-11,6
+11,7
@@
import matplotlib.pyplot as plt
import torch, torchvision
from torch import nn
import torch, torchvision
from torch import nn
+from torch.nn import functional as F
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@
-113,6
+114,9
@@
parser.add_argument('--data',
type = str, default = 'gaussian_mixture',
help = f'Toy data-set to use: {data_list}')
type = str, default = 'gaussian_mixture',
help = f'Toy data-set to use: {data_list}')
+parser.add_argument('--no_window',
+ action='store_true', default = False)
+
args = parser.parse_args()
if args.seed >= 0:
args = parser.parse_args()
if args.seed >= 0:
@@
-223,7
+227,7
@@
print(f'nb_parameters {sum([ p.numel() for p in model.parameters() ])}')
######################################################################
# Generate
######################################################################
# Generate
-def generate(size, 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():
with torch.no_grad():
@@
-280,72
+284,97
@@
if ema is not None: ema.copy_to_model()
model.eval()
model.eval()
-if train_input.dim() == 2:
+########################################
+# Nx1 -> histogram
+if train_input.dim() == 2 and train_input.size(1) == 1:
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
- # Nx1 -> histogram
- if train_input.size(1) == 1:
+ x = generate((10000, 1), T, alpha, alpha_bar, sigma,
+ model, train_mean, train_std)
- x = generate((10000, 1), 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)
- 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 = 'stepfilled', 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'diffusion_{args.data}.pdf'
+ print(f'saving {filename}')
+ fig.savefig(filename, bbox_inches='tight')
- # Nx2 -> scatter plot
- elif train_input.size(1) == 2:
+ if not args.no_window and hasattr(plt.get_current_fig_manager(), 'window'):
+ plt.get_current_fig_manager().window.setGeometry(2, 2, 2048, 768)
+ plt.show()
- x = generate((1000, 2), alpha, alpha_bar, sigma,
- model, train_mean, train_std)
+########################################
+# Nx2 -> scatter plot
+elif train_input.dim() == 2 and train_input.size(1) == 2:
- ax.set_xlim(-1.5, 1.5)
- ax.set_ylim(-1.5, 1.5)
- ax.set(aspect = 1)
- ax.spines.right.set_visible(False)
- ax.spines.top.set_visible(False)
+ fig = plt.figure()
+ ax = fig.add_subplot(1, 1, 1)
- d = x.detach().to('cpu').numpy()
- ax.scatter(d[:, 0], d[:, 1],
- s = 2.0, color = 'red', label = 'Synthesis')
+ 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.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.0
, 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')
- ax.legend(frameon = False, loc = 2)
+ 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)
filename = f'diffusion_{args.data}.pdf'
print(f'saving {filename}')
fig.savefig(filename, bbox_inches='tight')
filename = f'diffusion_{args.data}.pdf'
print(f'saving {filename}')
fig.savefig(filename, bbox_inches='tight')
- if 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()
plt.get_current_fig_manager().window.setGeometry(2, 2, 1024, 768)
plt.show()
+########################################
# NxCxHxW -> image
elif train_input.dim() == 4:
# NxCxHxW -> image
elif train_input.dim() == 4:
- x = generate((128,) + train_input.size()[1:], alpha, alpha_bar, sigma,
+ x = generate((128,) + train_input.size()[1:],
T,
alpha, alpha_bar, sigma,
model, train_mean, train_std)
model, train_mean, train_std)
- x = 1 - x.clamp(min = 0, max = 255) / 255
+
+ 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]
+
+ result = 1 - torch.cat((t, x), 2) / 255
filename = f'diffusion_{args.data}.png'
print(f'saving {filename}')
filename = f'diffusion_{args.data}.png'
print(f'saving {filename}')
- torchvision.utils.save_image(x, filename, nrow = 16, pad_value = 0.8)
+ torchvision.utils.save_image(result, filename)
+
+else:
+
+ print(f'cannot plot result of size {train_input.size()}')
######################################################################
######################################################################