Spaces:
Running
Running
| import matplotlib.pyplot as plt | |
| import torch | |
| import torchvision.utils as vutils | |
| from torch import nn | |
| def plot_example_grid( | |
| X, | |
| nrow=4, | |
| column_titles=None, | |
| cmap=None, | |
| normalize=True, | |
| renormalize_fn=None, | |
| title=None, | |
| figsize=None, | |
| save_path=None, | |
| dpi=80, | |
| ): | |
| """ | |
| X: tensor (B, C, H, W) | |
| nrow: number of images in each row | |
| column_titles: list of str - only if B is a multiple of nrow | |
| """ | |
| X = X.detach().cpu() | |
| if renormalize_fn: | |
| X = renormalize_fn(X) | |
| normalize = False | |
| grid = vutils.make_grid( | |
| X, nrow=nrow, normalize=normalize, scale_each=True, padding=2 | |
| ) | |
| npimg = grid.permute(1, 2, 0).numpy() | |
| H, W = npimg.shape[:2] | |
| img_h = H // (len(X) // nrow) | |
| img_w = W // nrow | |
| if figsize is not None: | |
| plt.figure(figsize=figsize, dpi=dpi) | |
| else: | |
| plt.figure(figsize=(min(20, 3.5 * nrow), 3.5 * (len(X) // nrow + 1)), dpi=dpi) | |
| plt.imshow(npimg, cmap=cmap or ("gray" if npimg.shape[-1] == 1 else None)) | |
| # Add column titles | |
| if column_titles is not None: | |
| for i, col_title in enumerate(column_titles): | |
| x_center = i * img_w + img_w / 2 | |
| plt.text(x_center, y=-5, s=col_title, fontsize=12, ha="center", va="bottom") | |
| if title: | |
| plt.title(title) | |
| plt.axis("off") | |
| plt.tight_layout() | |
| if save_path is not None: | |
| plt.savefig(save_path, bbox_inches="tight", pad_inches=0) | |
| else: | |
| plt.show() | |
| def plot_function( | |
| f, | |
| x_range=(-5, 5), | |
| num_points=1000, | |
| title="Function plot", | |
| xlabel="x", | |
| ylabel="f(x)", | |
| dpi=80, | |
| ): | |
| x = torch.linspace(x_range[0], x_range[1], num_points) | |
| y = f(x) | |
| y = y.detach().cpu().numpy() | |
| x = x.detach().cpu().numpy() | |
| plt.figure(figsize=(8, 4), dpi=dpi) | |
| plt.plot(x, y, label="f(x)") | |
| plt.axhline(0, color="black", linewidth=1, linestyle="--") # y=0 axis | |
| plt.title(title) | |
| plt.xlabel(xlabel) | |
| plt.ylabel(ylabel) | |
| plt.grid(True) | |
| plt.legend() | |
| plt.tight_layout() | |
| plt.show() | |
| def maxpool2d_param_extractor(child): | |
| return { | |
| "kernel_size": child.kernel_size, | |
| "stride": child.stride, | |
| "padding": child.padding, | |
| "dilation": child.dilation, | |
| "return_indices": child.return_indices, | |
| "ceil_mode": child.ceil_mode, | |
| } | |
| def replace_module_with_custom_( | |
| module, custom_cls, original_cls=None, by_name=None, param_extractor=None | |
| ): | |
| for name, child in module.named_children(): | |
| if (by_name is not None and name == by_name) or ( | |
| original_cls is not None and isinstance(child, original_cls) | |
| ): | |
| params = param_extractor(child) if param_extractor is not None else {} | |
| setattr(module, name, custom_cls(**params)) | |
| else: | |
| replace_module_with_custom_( | |
| child, | |
| custom_cls, | |
| original_cls=original_cls, | |
| by_name=by_name, | |
| param_extractor=param_extractor, | |
| ) | |
| def show_images(images, adv_images, k=5): | |
| uimages = images.unflatten(0, (5, 1)) | |
| uadv_images = adv_images.unflatten(0, (5, k)) | |
| udiff = uadv_images - uimages | |
| show_adv = torch.cat([uimages, uadv_images], dim=1).flatten(0, 1) | |
| show_diff = torch.cat([uimages, udiff], dim=1).flatten(0, 1) | |
| return show_adv, show_diff | |