Spaces:
Running on Zero
Running on Zero
File size: 3,340 Bytes
8861caf 2003ad1 8861caf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | import base64
import io
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from torch import optim
class FeatureVisualizer:
"""
Uses Gradient Ascent to hallucinate an input image that maximally activates a target neuron.
This version uses its OWN hook that preserves gradients (does NOT detach).
"""
def __init__(self, model: torch.nn.Module):
self.model = model
def generate_synthetic_image(
self,
layer_name: str,
channel_idx: int,
steps: int = 200,
lr: float = 0.05,
device: str = "cpu"
) -> str:
self.model.to(device)
self.model.eval()
# We need our own hook that does NOT detach (so gradients flow)
target_layer = None
for name, module in self.model.named_modules():
if name == layer_name:
target_layer = module
break
if target_layer is None:
raise ValueError(f"Layer {layer_name} not found")
captured_activation = {}
def grad_hook(module, input, output):
# Store WITHOUT detaching so backward() can flow through
captured_activation['value'] = output
handle = target_layer.register_forward_hook(grad_hook)
# ResNet accepts smaller spatial inputs. Starting at 96px makes the
# interactive feature-visualization endpoint finish on CPU Spaces;
# the result is enlarged for display below.
image_tensor = torch.randn((1, 3, 96, 96), device=device) * 0.01
image_tensor = image_tensor.requires_grad_(True)
optimizer = optim.Adam([image_tensor], lr=lr, weight_decay=1e-6)
for step in range(steps):
optimizer.zero_grad()
_ = self.model(image_tensor)
acts = captured_activation['value']
if len(acts.shape) == 4:
target_activation = acts[0, channel_idx, :, :].mean()
else:
target_activation = acts[0, channel_idx]
loss = -target_activation
loss.backward()
optimizer.step()
# Blur every 10 steps to reduce high-frequency noise
if step % 10 == 0:
with torch.no_grad():
image_tensor.data = F.avg_pool2d(
image_tensor.data, kernel_size=3, stride=1, padding=1
)
handle.remove()
# Post-process into viewable image
img_data = image_tensor.detach().cpu().squeeze().numpy()
img_data = np.transpose(img_data, (1, 2, 0))
for c in range(3):
ch = img_data[:, :, c]
ch_min, ch_max = ch.min(), ch.max()
if ch_max - ch_min > 1e-5:
img_data[:, :, c] = (ch - ch_min) / (ch_max - ch_min)
else:
img_data[:, :, c] = 0.5
img_uint8 = (np.clip(img_data, 0, 1) * 255).astype(np.uint8)
pil_img = Image.fromarray(img_uint8).resize((512, 512), Image.LANCZOS)
buffered = io.BytesIO()
pil_img.save(buffered, format="PNG")
return base64.b64encode(buffered.getvalue()).decode("utf-8")
|