Giter VIP home page Giter VIP logo

torch-dreams's Introduction

Torch-Dreams

Making neural networks more interpretable, for research and art.

Open In Colab build codecov

pip install torch-dreams 

Contents:

Minimal example

Make sure you also check out the quick start colab notebook

import matplotlib.pyplot as plt
import torchvision.models as models
from torch_dreams import Dreamer

model = models.inception_v3(pretrained=True)
dreamy_boi = Dreamer(model, device = 'cuda')

image_param = dreamy_boi.render(
    layers = [model.Mixed_5b],
)

plt.imshow(image_param)
plt.show()

Not so minimal example

model = models.inception_v3(pretrained=True)
dreamy_boi = Dreamer(model, device = 'cuda', quiet = False)

image_param = dreamy_boi.render(
    layers = [model.Mixed_5b],
    width = 256,
    height = 256,
    iters = 150,
    lr = 9e-3,
    rotate_degrees = 15,
    scale_max = 1.2,
    scale_min =  0.5,
    translate_x = 0.2,
    translate_y = 0.2,
    custom_func = None,
    weight_decay = 1e-2,
    grad_clip = 1.,
)

plt.imshow(image_param)
plt.show()

Visualizing individual channels with custom_func

model = models.inception_v3(pretrained=True)
dreamy_boi = Dreamer(model, device = 'cuda')

layers_to_use = [model.Mixed_6b.branch1x1.conv]

def make_custom_func(layer_number = 0, channel_number= 0): 
    def custom_func(layer_outputs):
        loss = layer_outputs[layer_number][:, channel_number].mean()
        return -loss
    return custom_func

my_custom_func = make_custom_func(layer_number= 0, channel_number = 119)

image_param = dreamy_boi.render(
    layers = layers_to_use,
    custom_func = my_custom_func,
)
plt.imshow(image_param)
plt.show()

Batched generation for large scale experiments

The BatchedAutoImageParam paired with the BatchedObjective can be used to generate multiple feature visualizations in parallel. This takes up more memory based on the batch size, but is also faster than generating one visualization at a time.

from torch_dreams import Dreamer
import torchvision.models as models
from torch_dreams.batched_objective import BatchedObjective
from torch_dreams.batched_image_param import BatchedAutoImageParam

model = models.inception_v3(pretrained=True)
dreamy_boi = Dreamer(model, device="cuda")

## specify list of neuron indices to visualize
batch_neuron_indices = [i for i in range(10,20)]

## set up a batch of trainable image parameters
bap = BatchedAutoImageParam(
    batch_size=len(batch_neuron_indices), 
    width=256, 
    height=256, 
    standard_deviation=0.01
)

## objective generator for each neuron
def make_custom_func(layer_number=0, channel_number=0):
    def custom_func(layer_outputs):
        loss = layer_outputs[layer_number][:, channel_number].norm()
        return -loss

    return custom_func

## prepare objective functions for each neuron index
batched_objective = BatchedObjective(
    objectives=[make_custom_func(channel_number=i) for i in batch_neuron_indices]
)

## render activation maximization signals
result_batch = dreamy_boi.render(
    layers=[model.Mixed_5b],
    image_parameter=bap,
    iters=120,
    custom_func=batched_objective,
)

## save results in a folder
for i in batch_neuron_indices:
    result_batch[batch_neuron_indices.index(i)].save(f"results/{i}.jpg")

Caricatures

Caricatures create a new image that has a similar but more extreme activation pattern to the input image at a given layer (or multiple layers at a time). It's inspired from this issue

In this case, let's use googlenet

model = models.googlenet(pretrained = True)
dreamy_boi = Dreamer(model = model, quiet= False, device= 'cuda')

image_param = dreamy_boi.caricature(
    input_tensor = image_tensor, 
    layers = [model.inception4c],   ## feel free to append more layers for more interesting caricatures 
    power= 1.2,                     ## higher -> more "exaggerated" features
)

plt.imshow(image_param)
plt.show()

Visualize features from multiple models on a single image parameter

First, let's pick 2 models and specify which layers we'd want to work with

from torch_dreams.model_bunch import ModelBunch

bunch = ModelBunch(
    model_dict = {
        'inception': models.inception_v3(pretrained=True).eval(),
        'resnet':    models.resnet18(pretrained= True).eval()
    }
)

layers_to_use = [
            bunch.model_dict['inception'].Mixed_6a,
            bunch.model_dict['resnet'].layer2[0].conv1
        ]

dreamy_boi = Dreamer(model = bunch, quiet= False, device= 'cuda')

Then define a custom_func which determines which exact activations of the models we have to optimize

def custom_func(layer_outputs):
    loss =   layer_outputs[0].mean()*2.0 + layer_outputs[1][:, 89].mean() 
    return -loss

Run the optimization

image_param = dreamy_boi.render(
    layers = layers_to_use,
    custom_func= custom_func,
    iters= 100
)

plt.imshow(image_param)
plt.show()

Using custom transforms:

import torchvision.transforms as transforms

model = models.inception_v3(pretrained=True)
dreamy_boi = Dreamer(model,  device = 'cuda', quiet =  False)

my_transforms = transforms.Compose([
    transforms.RandomAffine(degrees = 10, translate = (0.5,0.5)),
    transforms.RandomHorizontalFlip(p = 0.3)
])

dreamy_boi.set_custom_transforms(transforms = my_transforms)

image_param = dreamy_boi.render(
    layers = [model.Mixed_5b],
)

plt.imshow(image_param)
plt.show()

You can also use outputs of one render() as the input of another to create feedback loops.

import matplotlib.pyplot as plt
import torchvision.models as models
from torch_dreams import Dreamer

model = models.inception_v3(pretrained=True)
dreamy_boi = Dreamer(model,  device = 'cuda', quiet =  False)

image_param = dreamy_boi.render(
    layers = [model.Mixed_6c],
)

image_param = dreamy_boi.render(
    image_parameter= image_param,
    layers = [model.Mixed_5b],
    iters = 20
)

plt.imshow(image_param)
plt.show()

Using custom images

Note that you might have to use smaller values for certain hyperparameters like lr and grad_clip.

from torch_dreams.custom_image_param import CustomImageParam
param = CustomImageParam(image = 'images/sample_small.jpg', device= 'cuda')  ## image could either be a filename or a torch.tensor of shape NCHW

image_param = dreamy_boi.render(
    image_parameter= param,
    layers = [model.Mixed_6c],
    lr = 2e-4,
    grad_clip = 0.1,
    weight_decay= 1e-1,
    iters = 120
)

Working on models with different image normalizations

torch-dreams generally works with models trained on images normalized with imagenet mean and std, but that can be easily overriden to support any other normalization. For example, if you have a model with mean = [0.5, 0.5, 0.5] and std = [0.5, 0.5, 0.5]:

t = torchvision.transforms.Normalize(
                mean = [0.5, 0.5, 0.5],
                std =  [0.5, 0.5, 0.5]
            )

dreamy_boi.set_custom_normalization(normalization_transform = t) ## normalization_transform could be any instance of torch.nn.Module

Masked Image parameters

Can be used to optimize only certain parts of the image using a mask whose values are clipped between [0,1].

Here's an example with a vertical gradient

from torch_dreams.masked_image_param import MaskedImageParam

mask = torch.ones(1,1,512,512)

for i in range(0, 512, 1):  ## vertical gradient
    mask[:,:,i,:] = (i/512)

param = MaskedImageParam(
    image= 'images/sample_small.jpg',  ## optional
    mask_tensor = mask,
    device = 'cuda'
)

param = dreamy_boi.render(
    layers = [model.inception4c],
    image_parameter= param,
    lr = 1e-4,
    grad_clip= 0.1,
    weight_decay= 1e-1,
    iters= 200,
)

param.save('masked_param_output.jpg')

It's also possible to update the mask on the fly with param.update_mask(some_mask)

param.update_mask(mask = torch.flip(mask, dims = (2,)))

param = dreamy_boi.render(
    layers = [model.inception4a],
    image_parameter= param,
    lr = 1e-4,
    grad_clip= 0.1,
    weight_decay= 1e-1,
    iters= 200,
)

param.save('masked_param_output_2.jpg')

Other conveniences

The following methods are handy for an auto_image_param instance:

  1. Saving outputs as images:
image_param.save('output.jpg')
  1. Torch Tensor of dimensions (height, width, color_channels)
torch_image = image_param.to_hwc_tensor(device = 'cpu')
  1. Torch Tensor of dimensions (color_channels, height, width)
torch_image_chw = image_param.to_chw_tensor(device = 'cpu')
  1. Displaying outputs on matplotlib.
plt.imshow(image_param)
plt.show()
  1. For instances of custom_image_param, you can set any NCHW tensor as the image parameter:
image_tensor = image_param.to_nchw_tensor()

## do some stuff with image_tensor
t = transforms.Compose([
    transforms.RandomRotation(5)
])
transformed_image_tensor = t(image_tensor) 

image_param.set_param(tensor = transformed_image_tensor)

Args for render()

  • layers (iterable): List of the layers of model(s)'s layers to work on. [model.layer1, model.layer2...]

  • image_parameter (auto_image_param, optional): Instance of torch_dreams.auto_image_param.auto_image_param

  • width (int, optional): Width of image to be optimized

  • height (int, optional): Height of image to be optimized

  • iters (int, optional): Number of iterations, higher -> stronger visualization

  • lr (float, optional): Learning rate

  • rotate_degrees (int, optional): Max rotation in default transforms

  • scale_max (float, optional): Max image size factor. Defaults to 1.1.

  • scale_min (float, optional): Minimum image size factor. Defaults to 0.5.

  • translate_x (float, optional): Maximum translation factor in x direction

  • translate_y (float, optional): Maximum translation factor in y direction

  • custom_func (function, optional): Can be used to define custom optimiziation conditions to render(). Defaults to None.

  • weight_decay (float, optional): Weight decay for default optimizer. Helps prevent high frequency noise. Defaults to 0.

  • grad_clip (float, optional): Maximum value of the norm of gradient. Defaults to 1.

Args for Dreamer.__init__()

  • model (nn.Module or torch_dreams.model_bunch.Modelbunch): Almost any PyTorch model which was trained on imagenet mean and std, and supports variable sized images as inputs. You can pass multiple models into this argument as a torch_dreams.model_bunch.Modelbunch instance.
  • quiet (bool): Set to True if you want to disable any progress bars
  • device (str): cuda or cpu depending on your runtime

Development

  1. Clone the repo and navigate into the folder
git clone [email protected]:Mayukhdeb/torch-dreams.git
cd torch-dreams/
  1. Install dependencies
pip install -r requirements.txt
  1. Install torch-dreams as an editable module
python3 setup.py develop

Citation

@misc{mayukhdebtorchdreams,
  title={Feature Visualization library for PyTorch},
  author={Mayukh Deb},
  year={2021},
  publisher={GitHub},
  howpublished={\url{https://github.com/Mayukhdeb/torch-dreams}},
}

Acknowledgements

Recommended Reading

torch-dreams's People

Contributors

genekogan avatar imgbotapp avatar mayukhdeb avatar sushmanthreddy avatar

Stargazers

 avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar

Watchers

 avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar

torch-dreams's Issues

get rid of opencv

it's a relatively big dependency in terms of size, but I'm barely getting any use out of it.

Steps:

  1. find out where opencv is used
  2. find viable replacements of those opencv operations with numpy, PIL or ImageIO
  3. replace them one at a time
  4. pass all tests

add support for n channel image parameters

something like:

from torch_dreams.n_channel_param import NChannelParam

param = NChannelParam(
    num_channels = 2, 
    mean = [0.5,0.5], 
    std = [0.5, 0.5],
    device = 'cuda'
)

out = dreamy_boi.render(
    layers = [model.layer3],
    image_parameter = param
)

add support for noisegrad

The paper can be found here: https://arxiv.org/pdf/2106.10185.pdf
The core idea is to add noise to the weights on each iteration, in proportion to the weights variance in each layer.

It can be added in as an argument in the render() method:

image_param = dreamy_boi.render(
    layers = [model.Mixed_5b],
    use_noisegrad = True
)

confusion about visualizing individual channels

As a freshman in this area, I really appreciate your repo for letting me know how feature visualization works in practice. Thanks!!

In the channel visualization case, the channel index in the example is 119, however, the colab tutorial states the 7th channel in Mixed_6b.branch1x1 is optimized. In my understanding, I think the index of the optimized channel should equal the value of channel_number (119th channel in this case). Or is there something I misunderstood?

drawing bigger features

thanks for repo, quite fun and cool!
it seems, the features on example images are drawn bigger, than it happens when using this library. is there any trick for that (besides external upscaling)?

masked image params

Masked image params can be thought of as custom image parameters which can be updated only on certain parts of the image.

It can work like:

from torch_dreams.masked_image_param import MaskedImageParam

param = MaskedImageParam(
    image = 'image.jpg',  ## or hwc np array 
    mask = mask_tensor_or_hwc_numpy,
    device = 'cuda'
)
    
dreamy_boi.render(
    param = param, 
    layers = [layer]
)

On each forward pass, the image gets multiplied by a mask in the spatial domain.

build ImageParam class for handling image updates

It could be something like:

class image_param():
    def __init__(self, image_tensor):
        self.tensor = image_tensor
        self.tensor.requires_grad = True

    def get_optimizer(self, lr, optimizer = None, momentum = 0., weight_decay = 0.):
        if optimizer is None:
            self.optimizer = optim.SGD([self.tensor], lr = lr, momentum=momentum, weight_decay=weight_decay)
        else:
            self.optimizer = optimizer([self.tensor], lr = lr, weight_decay = weight_decay)

    def more_stuff(...)

revamp custom_func args

all the layer outputs can be stored in a dict instead of a list.

it could look like:

layer_outputs = {
    '__input_image__': input_image_tensor,
    'conv3.branch1': some_tensor,
    'fc': another_tensor
}

and custom_func might look like:

def custom_func(outs):
    loss = outs['conv3.branch1'][:100] + outs['fc'] + torch.var(outs['__input_image__'])
    return -loss

broken pipeline on torch v1.8.0+cu111

the latest non torch-dreams functions do not run on torch v1.8.0+cu111 (which is on google colab), but runs on torch v1.7.0

This can be traced back to torch.irfft on v1.7.0. Which needs to be replaced by torch.fft,irfftn on v1.8.0+cu111

static caricatures

The current version of torch_dreams.dreamer.caricature() dynamically updates the target layer outputs in each step according to the transforms.

This won't work well for experiments with adversarial examples, since their outcomes are dependent on the position of each pixel. Hence a solution would be to generate the layer outputs at first and keep use them at each iteration as the target.

This can be integrate into the interface as:

image_param = dreamy_boi.caricature(
    input_tensor = image_tensor, 
    layers = [model.some_layer],  
    power= 1.2,   
    static = True                
)

Something like:
adversarial_caricature

v3.0

Todo:

  • use proper naming styles as seen here (breaking change)
  • add support for batched generation (would lead to a big speedup on Dora)
  • more tests!
  • update documentation (include steps for local dev + running tests)

extra tools for torch-dreams

This might contain stuff which generally require a large amount of boilerplate in notebooks/scripts, but are used often in code

Something like:

from torch_dreams.recipes.losses import SnapshotLoss

Obviously inspired from lucid.recipes

and:

from torch_dreams.tools import get_video_writer
writer = get_video_writer(filename = 'save.mp4', framerate = 60)
## some stuff
writer.append_data(some_image_param.to_hwc_tensor().numpy())

writer.close()

move to torch.optim for optimization

Currently, the optimization is done using a simple implementation which goes as:

image_tensor.data = image_tensor.data + lr *(gradients_tensor.data /grad_norm)

to something like

optim = optim.some_optimizer(image_tensor.parameters(), lr= config["learning_rate"], momentum=0.9)
loss.backward()

Have a loot at the shampoo optimizer, it can be used for "preconditioning"

Two advantages:

  • Cleaner code
  • (maybe) better performance

add support for models on any image normalization

For example, this transform converts an imagenet transformed image tensor to 0, 255, this can be generalized.

inv = transforms.Compose([ transforms.Normalize(mean = [ 0., 0., 0. ],
                                                     std = [ 1/0.229, 1/0.224, 1/0.225 ]),
                                transforms.Normalize(mean = [ -0.485, -0.456, -0.406 ],
                                                     std = [ 1., 1., 1. ]),
                               ])   

class InverseTransform(nn.Module):
    def __init__(self, inverse_transform):
        super().__init__()
        self.inverse_transform = inverse_transform
        self.scale = 255

    def forward(self, x):
        x=  self.inverse_transform(x)
        x = x * self.scale
        # print(x.max(), x.min(), x.mean())
        return x

torch_dreams.auto: add support for multiple models

This might be possible with torch_dreams.auto.model_bunch.ModelBunch as:

from torch_dreams.auto.model_bunch import ModelBunch

model1  = models.inception_v3(pretrained=True).eval()
model2  = models.resnet18(pretrained= True).eval() 

bunch = ModelBunch(
    model_dict = {
        'inception': model1,
        'resnet':  model2
    }
)

layers_to_use = [
    bunch.model_dict['inception'].Mixed_6a,
    bunch.model_dict['resnet'].layer2[0].conv1
]

dreamy_boi = dreamer(model = bunch, quiet= False, device= 'cuda')

custom_func should also include the input image tensor

This'll enable us to feed the input image into a a discriminator for more realistic examples

something like:

def custom_func(input_image, layer_outputs):
     loss = discriminator(mean) + layer_outputs[0].mean()
     return -loss

Reproducibility through config dictionaries

The current __init__ function for torch_dreams.dreamer is very basic, it would be better if all the octave sizes and random roll/rotate values are generated before the algorithm starts. These variables should also be accessible with something like:

self.config = {
    "octave_scale": 1.3,
    "num_octaves": 7,
    "iterations": 20,
    "lr": 0.03,
    "gradient_smoothing_coeff": 1.5,
    "gradient_smoothing_kernel_size": 9,
    "rotations": np.array([
                            [0.1,0.3,..],  ## shape would be (num_octaves, iterations)
                            .
                            [0.2,0.1...]
                        ]),
    
    "rolls":np.array([
                    [0.1,0.3,..],  ## shape would be (num_octaves, iterations)
                    .
                    [0.2,0.1...]
                ])
    
}

Then the results could be reproduced by:

out = dreamer.deep_dream_on_config(config)

octave_utils.py: UnboundLocalError when setting either of smoothing args to None

If we use

config = {
    ...other values...
    "gradient_smoothing_coeff": None,
    "gradient_smoothing_kernel_size": None,
}

and then run dreamer.deep_dream_with_masks(config), it throws the following traceback:

Traceback (most recent call last):
  File "test_gradient_masks.py", line 50, in <module>
    out_single_conv_a = dreamy_boi.deep_dream_with_masks(config)
  File "/media/mayukh/Data/storage/repositories/repos/torch-dreams/torch_dreams/dreamer.py", line 150, in deep_dream_with_masks
    image_np = self.dream_on_octave_with_masks(
  File "/media/mayukh/Data/storage/repositories/repos/torch-dreams/torch_dreams/octave_utils.py", line 78, in dream_on_octave_with_masks
    image_tensor.data = image_tensor.data + lr *(gradients_tensor.data /g_norm) ## can confirm this is still on the GPU if you have one
UnboundLocalError: local variable 'g_norm' referenced before assignment

RuntimeError randomly when static = True on caricatures

This occurs randomly, and the most likely reason this snippet:

The reason is that if the shape of x is the same as the shape of y, we don't use resize and hence also dont run the .detach() on y

if x.shape != y.shape:
"""
if their shapes are not equal (likely due to using static caricatures), then resize the target accordingly
"""
y = resize_4d_tensor_by_size(y.unsqueeze(0), height = x.shape[-2], width = x.shape[-1] ).squeeze(0).detach()
numerator = (x*y).sum()

Traceback:

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-174-9f4719af0a34> in <module>()
      9         power= 1.2,                     ## higher -> more "exaggerated" features
     10         static = True,
---> 11         iters = 200
     12     )
     13 

2 frames
/usr/local/lib/python3.7/dist-packages/torch_dreams/dreamer.py in caricature(self, input_tensor, layers, power, image_parameter, iters, lr, rotate_degrees, scale_max, scale_min, translate_x, translate_y, weight_decay, grad_clip, static)
    246             loss = caricature_loss.forward(layer_outputs = layer_outputs, ideal_layer_outputs = ideal_layer_outputs)
    247 
--> 248             loss.backward()
    249             image_parameter.clip_grads(grad_clip= grad_clip)
    250             image_parameter.optimizer.step()

/usr/local/lib/python3.7/dist-packages/torch/tensor.py in backward(self, gradient, retain_graph, create_graph, inputs)
    243                 create_graph=create_graph,
    244                 inputs=inputs)
--> 245         torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
    246 
    247     def register_hook(self, hook):

/usr/local/lib/python3.7/dist-packages/torch/autograd/__init__.py in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)
    145     Variable._execution_engine.run_backward(
    146         tensors, grad_tensors_, retain_graph, create_graph, inputs,
--> 147         allow_unreachable=True, accumulate_grad=True)  # allow_unreachable flag
    148 
    149 

RuntimeError: Trying to backward through the graph a second time, but the saved intermediate results have already been freed. Specify retain_graph=True when calling .backward() or autograd.grad() the first time.

layers_to_use should be a dict, not a list

instead of using:

layers_to_use = [model.layer1, model.layer2]

and then accessing them in the custom function as:

layer_outputs[0] and layer_outputs[1]

we can define layers_to_use as a dict which makes custom_func much more intuitive

layers_to_use = {
    'layer1': model.layer1, 
    'layer2': model.layer2
}

and then accessing them in the custom function as:

layer_outputs['layer2'] and layer_outputs['layer1']

Support for latest pytorch 1.9

according to the requirements.txt, torch>=1.8.0 is required. However the code does not run when the latest version pytorch 1.9 is used. The following line will raise an exception in that case:

if torch.__version__[:3] == '1.8':

If only 1.8.0 works, then the requirements should be updated to that specific version, otherwise the version check in the code should be changed.

Recommend Projects

  • React photo React

    A declarative, efficient, and flexible JavaScript library for building user interfaces.

  • Vue.js photo Vue.js

    ๐Ÿ–– Vue.js is a progressive, incrementally-adoptable JavaScript framework for building UI on the web.

  • Typescript photo Typescript

    TypeScript is a superset of JavaScript that compiles to clean JavaScript output.

  • TensorFlow photo TensorFlow

    An Open Source Machine Learning Framework for Everyone

  • Django photo Django

    The Web framework for perfectionists with deadlines.

  • D3 photo D3

    Bring data to life with SVG, Canvas and HTML. ๐Ÿ“Š๐Ÿ“ˆ๐ŸŽ‰

Recommend Topics

  • javascript

    JavaScript (JS) is a lightweight interpreted programming language with first-class functions.

  • web

    Some thing interesting about web. New door for the world.

  • server

    A server is a program made to process requests and deliver data to clients.

  • Machine learning

    Machine learning is a way of modeling and interpreting data that allows a piece of software to respond intelligently.

  • Game

    Some thing interesting about game, make everyone happy.

Recommend Org

  • Facebook photo Facebook

    We are working to build community through open source technology. NB: members must have two-factor auth.

  • Microsoft photo Microsoft

    Open source projects and samples from Microsoft.

  • Google photo Google

    Google โค๏ธ Open Source for everyone.

  • D3 photo D3

    Data-Driven Documents codes.