ProGamerGov/image_cppn.py

## image_cppn.py
# tensorflow/lucid CPPN (X,Y) --> (R,G,B) Differentiable Image Parameterization in PyTorch
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
from inception5h import Inception5h
from PIL import Image


def main():
    # Setup params
    image_size=(224,224)
    iterations = 512
    lr = 0.005
    use_device = 'cuda:0'

    # Uncomment to use seed
    #set_seed(876)

    # Target layer and channel
    layer_name='mixed4b_3x3_pre_relu_conv'
    channel=77

    # Load InceptionV1
    cnn = Inception5h()
    cnn.load_state_dict(torch.load('inception5h.pth'))
    cnn.add_layers() # Add layers not found in model file
    net = cnn.to(use_device).eval()
    for param in net.parameters():
        param.requires_grad = False

    # Uncomment to print all hookable layers
    #print(get_hookable_layers(cnn))

    # Create instance of CPPN
    img_cppn = CPPN_Conv(size=image_size, num_channels=24, num_layers=8, use_device=use_device, normalize=False)
    print('CPPN Params Count', sum(p.numel() for p in img_cppn.parameters() if p.requires_grad))

    loss_modules = register_hook_module(net, layer_name, channel)

    # Create 224x224 image
    output_tensor, img_cppn = dream_cppn(net, img_cppn, iterations, lr, loss_modules, use_device)
    simple_deprocess(output_tensor, name='out.png')

    # Create 720x720 image using multiscale generation
    image_size=(720,720)
    img_cppn.remake_input(image_size, use_device=use_device)
    output_tensor, img_cppn = dream_cppn(net, img_cppn, iterations, lr, loss_modules, use_device)
    simple_deprocess(output_tensor, name='out.png')


# Simple Deprocess
def simple_deprocess(output_tensor, name):
    output_tensor = output_tensor.squeeze(0).cpu() /255
    output_tensor.clamp_(0, 1)
    Image2PIL = transforms.ToPILImage()
    image = Image2PIL(output_tensor.cpu())
    image.save(name)


# Function to maximize CNN activations
def dream_cppn(net, img_cppn, iterations, lr, loss_modules, use_device):
    # Setup optimizer to optimize CPPN instance
    optimizer = torch.optim.Adam(img_cppn.parameters(), lr=lr)

    # Training loop
    for i in range(iterations):
        optimizer.zero_grad()
        img = img_cppn() * 255 # Create RGB image with CPPN
        net(img) # Create loss values
        loss = 0
        for mod in loss_modules: # Collect loss values
            loss += mod.loss
        loss.backward()

        # Uncomment to save iterations
        #if i % 25 == 0:
        #     simple_deprocess(img.detach(), 'out_'+str(i)+'.png')

        print('Iteration', str(i+1), 'Loss', str(loss.item()))

        optimizer.step()
    img = img_cppn() * 255
    return img, img_cppn


# Activation function for CPPN
class CompositeActivation(nn.Module):
    def forward(self, input):
        input = torch.atan(input)
        return torch.cat([input / 0.67, (input * input) / 0.6], 1)


# Compositional pattern-producing network (CPPN) with Conv2d layers
class CPPN_Conv(nn.Module):

    def __init__(self, size=(224, 224), num_channels=24, num_layers=8, activ_func=CompositeActivation(), use_device='cpu', normalize=False):
        super(CPPN_Conv, self).__init__()
        self.input_size = size
        self.n_channels = num_channels
        self.net = self.create_net(num_channels, num_layers, activ_func, use_device, normalize)
        self.cppn_input = self.create_input(use_device)

    # Create CPPN (X,Y) --> (R,G,B)
    def create_net(self, num_channels, num_layers, activ_func, use_device, instance_norm, affine=True, bias=True):
        net = nn.Sequential()
        net.add_module(str(len(net)), nn.Conv2d(in_channels=2, out_channels=num_channels, kernel_size=1, bias=bias))
        if instance_norm:
            net.add_module(str(len(net)), nn.InstanceNorm2d(num_channels, affine=affine))
        net.add_module(str(len(net)), activ_func)
        for l in range(num_layers - 1):
            net.add_module(str(len(net)), nn.Conv2d(in_channels=num_channels*2, out_channels=num_channels, kernel_size=1, bias=bias))
            if instance_norm:
                 net.add_module(str(len(net)), nn.InstanceNorm2d(num_channels, affine=affine))
            net.add_module(str(len(net)), activ_func)
        net.add_module(str(len(net)), nn.Conv2d(in_channels=num_channels*2, out_channels=3, kernel_size=1, bias=bias))
        net.add_module(str(len(net)), nn.Sigmoid())
        net.apply(self.cppn_normal)
        return net.to(use_device)

    # Create X,Y input for CPPN
    def create_input(self, use_device):
        if type(self.input_size) is not tuple and type(self.input_size) is not list:
            self.input_size = (self.input_size, self.input_size)
        w = torch.arange(0, self.input_size[1]).to(use_device)
        h = torch.arange(0, self.input_size[0]).to(use_device)
        w_exp = w.unsqueeze(1).expand((self.input_size[1], self.input_size[0])).true_divide(self.input_size[0]) - 0.5
        h_exp = h.unsqueeze(0).expand((self.input_size[1], self.input_size[0])).true_divide(self.input_size[1]) - 0.5
        return torch.stack((w_exp, h_exp), -1).permute(2,1,0).unsqueeze(0)

    # Replace input with different sized input
    def remake_input(self, image_size, use_device):
        self.input_size = image_size
        self.cppn_input = self.create_input(use_device)

    # Normalize Conv2d weights
    def cppn_normal(self, l):
        if type(l) == nn.Conv2d:
            l.weight.data.normal_(0, (1.0/self.n_channels)**(1/2))
            if l.bias is not None:
                nn.init.zeros_(l.bias)

    def forward(self):
        return self.net(self.cppn_input)


# Create loss module and hook
def register_hook_module(net, layer_name, channel=-1):
    loss_module = SimpleDreamLossHook(channel)
    getattr(net, layer_name).register_forward_hook(loss_module)
    return [loss_module]


# Define a simple forward hook to collect DeepDream loss
class SimpleDreamLossHook(nn.Module):
    def __init__(self, channel=-1):
        super(SimpleDreamLossHook, self).__init__()
        self.get_loss = self.mean_loss
        self.channel = channel

    def mean_loss(self, input):
        return input.mean()

    def forward(self, module, input, output):
        if self.channel > -1:
            self.loss = -self.get_loss(output[:,self.channel,:,:])
        else:
            self.loss = -self.get_loss(output)


# Set global seeds to output reproducible
def set_seed(seed):
    import random
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic=True
    random.seed(seed)


# Get layers that can be hooked
def get_hookable_layers(cnn):
    return [l[0] for l in list(cnn.named_children())]


if __name__ == "__main__":
    main()

## inception5h.py
# Model donwload: https://github.com/ProGamerGov/pytorch-old-tensorflow-models/blob/master/inception5h.pth
import torch
import torch.nn as nn
import torch.nn.functional as F


class ReluLayer(nn.Module):
    def forward(self, input):
        return F.relu(input)

class CatLayer(nn.Module):
    def forward(self, input_list, dim=1):
        return torch.cat(input_list, dim)

class MaxPool2dLayer(nn.Module):
    def forward(self, input, kernel_size=(3, 3), stride=(1, 1), padding=0, ceil_mode=False):
        return F.max_pool2d(input, kernel_size, stride=stride, padding=padding, ceil_mode=ceil_mode)

class SoftMaxLayer(nn.Module):
    def forward(self, input, dim=1):
        return F.softmax(input, dim=dim)


class Inception5h(nn.Module):
    def __init__(self):
        super(Inception5h, self).__init__()
        self.conv2d0_pre_relu_conv = nn.Conv2d(in_channels=3, out_channels=64, kernel_size=(7, 7), stride=(2, 2), groups=1, bias=True)
        self.conv2d1_pre_relu_conv = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=(1, 1), stride=(1, 1), groups=1, bias=True)
        self.conv2d2_pre_relu_conv = nn.Conv2d(in_channels=64, out_channels=192, kernel_size=(3, 3), stride=(1, 1), groups=1, bias=True)
        self.mixed3a_1x1_pre_relu_conv = nn.Conv2d(in_channels=192, out_channels=64, kernel_size=(1, 1), stride=(1, 1), groups=1, bias=True)
        self.mixed3a_3x3_bottleneck_pre_relu_conv = nn.Conv2d(in_channels=192, out_channels=96, kernel_size=(1, 1), stride=(1, 1), groups=1, bias=True)
        self.mixed3a_5x5_bottleneck_pre_relu_conv = nn.Conv2d(in_channels=192, out_channels=16, kernel_size=(1, 1), stride=(1, 1), groups=1, bias=True)
        self.mixed3a_pool_reduce_pre_relu_conv = nn.Conv2d(in_channels=192, out_channels=32, kernel_size=(1, 1), stride=(1, 1), groups=1, bias=True)
        self.mixed3a_3x3_pre_relu_conv = nn.Conv2d(in_channels=96, out_channels=128, kernel_size=(3, 3), stride=(1, 1), groups=1, bias=True)
        self.mixed3a_5x5_pre_relu_conv = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=(5, 5), stride=(1, 1), groups=1, bias=True)
        self.mixed3b_1x1_pre_relu_conv = nn.Conv2d(in_channels=256, out_channels=128, kernel_size=(1, 1), stride=(1, 1), groups=1, bias=True)
        self.mixed3b_3x3_bottleneck_pre_relu_conv = nn.Conv2d(in_channels=256, out_channels=128, kernel_size=(1, 1), stride=(1, 1), groups=1, bias=True)
        self.mixed3b_5x5_bottleneck_pre_relu_conv = nn.Conv2d(in_channels=256, out_channels=32, kernel_size=(1, 1), stride=(1, 1), groups=1, bias=True)
        self.mixed3b_pool_reduce_pre_relu_conv = nn.Conv2d(in_channels=256, out_channels=64, kernel_size=(1, 1), stride=(1, 1), groups=1, bias=True)
        self.mixed3b_3x3_pre_relu_conv = nn.Conv2d(in_channels=128, out_channels=192, kernel_size=(3, 3), stride=(1, 1), groups=1, bias=True)
        self.mixed3b_5x5_pre_relu_conv = nn.Conv2d(in_channels=32, out_channels=96, kernel_size=(5, 5), stride=(1, 1), groups=1, bias=True)
        self.mixed4a_1x1_pre_relu_conv = nn.Conv2d(in_channels=480, out_channels=192, kernel_size=(1, 1), stride=(1, 1), groups=1, bias=True)
        self.mixed4a_3x3_bottleneck_pre_relu_conv = nn.Conv2d(in_channels=480, out_channels=96, kernel_size=(1, 1), stride=(1, 1), groups=1, bias=True)
        self.mixed4a_5x5_bottleneck_pre_relu_conv = nn.Conv2d(in_channels=480, out_channels=16, kernel_size=(1, 1), stride=(1, 1), groups=1, bias=True)
        self.mixed4a_pool_reduce_pre_relu_conv = nn.Conv2d(in_channels=480, out_channels=64, kernel_size=(1, 1), stride=(1, 1), groups=1, bias=True)
        self.mixed4a_3x3_pre_relu_conv = nn.Conv2d(in_channels=96, out_channels=204, kernel_size=(3, 3), stride=(1, 1), groups=1, bias=True)
        self.mixed4a_5x5_pre_relu_conv = nn.Conv2d(in_channels=16, out_channels=48, kernel_size=(5, 5), stride=(1, 1), groups=1, bias=True)
        self.mixed4b_1x1_pre_relu_conv = nn.Conv2d(in_channels=508, out_channels=160, kernel_size=(1, 1), stride=(1, 1), groups=1, bias=True)
        self.mixed4b_3x3_bottleneck_pre_relu_conv = nn.Conv2d(in_channels=508, out_channels=112, kernel_size=(1, 1), stride=(1, 1), groups=1, bias=True)
        self.mixed4b_5x5_bottleneck_pre_relu_conv = nn.Conv2d(in_channels=508, out_channels=24, kernel_size=(1, 1), stride=(1, 1), groups=1, bias=True)
        self.mixed4b_pool_reduce_pre_relu_conv = nn.Conv2d(in_channels=508, out_channels=64, kernel_size=(1, 1), stride=(1, 1), groups=1, bias=True)
        self.mixed4b_3x3_pre_relu_conv = nn.Conv2d(in_channels=112, out_channels=224, kernel_size=(3, 3), stride=(1, 1), groups=1, bias=True)
        self.mixed4b_5x5_pre_relu_conv = nn.Conv2d(in_channels=24, out_channels=64, kernel_size=(5, 5), stride=(1, 1), groups=1, bias=True)
        self.mixed4c_1x1_pre_relu_conv = nn.Conv2d(in_channels=512, out_channels=128, kernel_size=(1, 1), stride=(1, 1), groups=1, bias=True)
        self.mixed4c_3x3_bottleneck_pre_relu_conv = nn.Conv2d(in_channels=512, out_channels=128, kernel_size=(1, 1), stride=(1, 1), groups=1, bias=True)
        self.mixed4c_5x5_bottleneck_pre_relu_conv = nn.Conv2d(in_channels=512, out_channels=24, kernel_size=(1, 1), stride=(1, 1), groups=1, bias=True)
        self.mixed4c_pool_reduce_pre_relu_conv = nn.Conv2d(in_channels=512, out_channels=64, kernel_size=(1, 1), stride=(1, 1), groups=1, bias=True)
        self.mixed4c_3x3_pre_relu_conv = nn.Conv2d(in_channels=128, out_channels=256, kernel_size=(3, 3), stride=(1, 1), groups=1, bias=True)
        self.mixed4c_5x5_pre_relu_conv = nn.Conv2d(in_channels=24, out_channels=64, kernel_size=(5, 5), stride=(1, 1), groups=1, bias=True)
        self.mixed4d_1x1_pre_relu_conv = nn.Conv2d(in_channels=512, out_channels=112, kernel_size=(1, 1), stride=(1, 1), groups=1, bias=True)
        self.mixed4d_3x3_bottleneck_pre_relu_conv = nn.Conv2d(in_channels=512, out_channels=144, kernel_size=(1, 1), stride=(1, 1), groups=1, bias=True)
        self.mixed4d_5x5_bottleneck_pre_relu_conv = nn.Conv2d(in_channels=512, out_channels=32, kernel_size=(1, 1), stride=(1, 1), groups=1, bias=True)
        self.mixed4d_pool_reduce_pre_relu_conv = nn.Conv2d(in_channels=512, out_channels=64, kernel_size=(1, 1), stride=(1, 1), groups=1, bias=True)
        self.mixed4d_3x3_pre_relu_conv = nn.Conv2d(in_channels=144, out_channels=288, kernel_size=(3, 3), stride=(1, 1), groups=1, bias=True)
        self.mixed4d_5x5_pre_relu_conv = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=(5, 5), stride=(1, 1), groups=1, bias=True)
        self.mixed4e_1x1_pre_relu_conv = nn.Conv2d(in_channels=528, out_channels=256, kernel_size=(1, 1), stride=(1, 1), groups=1, bias=True)
        self.mixed4e_3x3_bottleneck_pre_relu_conv = nn.Conv2d(in_channels=528, out_channels=160, kernel_size=(1, 1), stride=(1, 1), groups=1, bias=True)
        self.mixed4e_5x5_bottleneck_pre_relu_conv = nn.Conv2d(in_channels=528, out_channels=32, kernel_size=(1, 1), stride=(1, 1), groups=1, bias=True)
        self.mixed4e_pool_reduce_pre_relu_conv = nn.Conv2d(in_channels=528, out_channels=128, kernel_size=(1, 1), stride=(1, 1), groups=1, bias=True)
        self.mixed4e_3x3_pre_relu_conv = nn.Conv2d(in_channels=160, out_channels=320, kernel_size=(3, 3), stride=(1, 1), groups=1, bias=True)
        self.mixed4e_5x5_pre_relu_conv = nn.Conv2d(in_channels=32, out_channels=128, kernel_size=(5, 5), stride=(1, 1), groups=1, bias=True)
        self.mixed5a_1x1_pre_relu_conv = nn.Conv2d(in_channels=832, out_channels=256, kernel_size=(1, 1), stride=(1, 1), groups=1, bias=True)
        self.mixed5a_3x3_bottleneck_pre_relu_conv = nn.Conv2d(in_channels=832, out_channels=160, kernel_size=(1, 1), stride=(1, 1), groups=1, bias=True)
        self.mixed5a_5x5_bottleneck_pre_relu_conv = nn.Conv2d(in_channels=832, out_channels=48, kernel_size=(1, 1), stride=(1, 1), groups=1, bias=True)
        self.mixed5a_pool_reduce_pre_relu_conv = nn.Conv2d(in_channels=832, out_channels=128, kernel_size=(1, 1), stride=(1, 1), groups=1, bias=True)
        self.mixed5a_3x3_pre_relu_conv = nn.Conv2d(in_channels=160, out_channels=320, kernel_size=(3, 3), stride=(1, 1), groups=1, bias=True)
        self.mixed5a_5x5_pre_relu_conv = nn.Conv2d(in_channels=48, out_channels=128, kernel_size=(5, 5), stride=(1, 1), groups=1, bias=True)
        self.mixed5b_1x1_pre_relu_conv = nn.Conv2d(in_channels=832, out_channels=384, kernel_size=(1, 1), stride=(1, 1), groups=1, bias=True)
        self.mixed5b_3x3_bottleneck_pre_relu_conv = nn.Conv2d(in_channels=832, out_channels=192, kernel_size=(1, 1), stride=(1, 1), groups=1, bias=True)
        self.mixed5b_5x5_bottleneck_pre_relu_conv = nn.Conv2d(in_channels=832, out_channels=48, kernel_size=(1, 1), stride=(1, 1), groups=1, bias=True)
        self.mixed5b_pool_reduce_pre_relu_conv = nn.Conv2d(in_channels=832, out_channels=128, kernel_size=(1, 1), stride=(1, 1), groups=1, bias=True)
        self.mixed5b_3x3_pre_relu_conv = nn.Conv2d(in_channels=192, out_channels=384, kernel_size=(3, 3), stride=(1, 1), groups=1, bias=True)
        self.mixed5b_5x5_pre_relu_conv = nn.Conv2d(in_channels=48, out_channels=128, kernel_size=(5, 5), stride=(1, 1), groups=1, bias=True)
        self.softmax2_pre_activation_matmul = nn.Linear(in_features = 1024, out_features = 1008, bias = True)

    def add_layers(self):
        self.conv2d0 = ReluLayer()
        self.maxpool0 = MaxPool2dLayer()
        self.conv2d1 = ReluLayer()
        self.conv2d2 = ReluLayer()
        self.maxpool1 = MaxPool2dLayer()
        self.mixed3a_pool = MaxPool2dLayer()
        self.mixed3a_1x1 = ReluLayer()
        self.mixed3a_3x3_bottleneck = ReluLayer()
        self.mixed3a_5x5_bottleneck = ReluLayer()
        self.mixed3a_pool_reduce = ReluLayer()
        self.mixed3a_3x3 = ReluLayer()
        self.mixed3a_5x5 = ReluLayer()
        self.mixed3a = CatLayer()
        self.mixed3b_pool = MaxPool2dLayer()
        self.mixed3b_1x1 = ReluLayer()
        self.mixed3b_3x3_bottleneck = ReluLayer()
        self.mixed3b_5x5_bottleneck = ReluLayer()
        self.mixed3b_pool_reduce = ReluLayer()
        self.mixed3b_3x3 = ReluLayer()
        self.mixed3b_5x5 = ReluLayer()
        self.mixed3b = CatLayer()
        self.maxpool4 = MaxPool2dLayer()
        self.mixed4a_pool = MaxPool2dLayer()
        self.mixed4a_1x1 = ReluLayer()
        self.mixed4a_3x3_bottleneck = ReluLayer()
        self.mixed4a_5x5_bottleneck = ReluLayer()
        self.mixed4a_pool_reduce = ReluLayer()
        self.mixed4a_3x3 = ReluLayer()
        self.mixed4a_5x5 = ReluLayer()
        self.mixed4a = CatLayer()
        self.mixed4b_pool = MaxPool2dLayer()
        self.mixed4b_1x1 = ReluLayer()
        self.mixed4b_3x3_bottleneck = ReluLayer()
        self.mixed4b_5x5_bottleneck = ReluLayer()
        self.mixed4b_pool_reduce = ReluLayer()
        self.mixed4b_3x3 = ReluLayer()
        self.mixed4b_5x5 = ReluLayer()
        self.mixed4b = CatLayer()
        self.mixed4c_pool = MaxPool2dLayer()
        self.mixed4c_1x1 = ReluLayer()
        self.mixed4c_3x3_bottleneck = ReluLayer()
        self.mixed4c_5x5_bottleneck = ReluLayer()
        self.mixed4c_pool_reduce = ReluLayer()
        self.mixed4c_3x3 = ReluLayer()
        self.mixed4c_5x5 = ReluLayer()
        self.mixed4c = CatLayer()
        self.mixed4d_pool = MaxPool2dLayer()
        self.mixed4d_1x1 = ReluLayer()
        self.mixed4d_3x3_bottleneck = ReluLayer()
        self.mixed4d_5x5_bottleneck = ReluLayer()
        self.mixed4d_pool_reduce = ReluLayer()
        self.mixed4d_3x3 = ReluLayer()
        self.mixed4d_5x5 = ReluLayer()
        self.mixed4d = CatLayer()
        self.mixed4e_pool = MaxPool2dLayer()
        self.mixed4e_1x1 = ReluLayer()
        self.mixed4e_3x3_bottleneck = ReluLayer()
        self.mixed4e_5x5_bottleneck = ReluLayer()
        self.mixed4e_pool_reduce = ReluLayer()
        self.mixed4e_3x3 = ReluLayer()
        self.mixed4e_5x5 = ReluLayer()
        self.mixed4e = CatLayer()
        self.maxpool10 = MaxPool2dLayer()
        self.mixed5a_pool = MaxPool2dLayer()
        self.mixed5a_1x1 = ReluLayer()
        self.mixed5a_3x3_bottleneck = ReluLayer()
        self.mixed5a_5x5_bottleneck = ReluLayer()
        self.mixed5a_pool_reduce = ReluLayer()
        self.mixed5a_3x3 = ReluLayer()
        self.mixed5a_5x5 = ReluLayer()
        self.mixed5a = CatLayer()
        self.mixed5b_pool = MaxPool2dLayer()
        self.mixed5b_1x1 = ReluLayer()
        self.mixed5b_3x3_bottleneck = ReluLayer()
        self.mixed5b_5x5_bottleneck = ReluLayer()
        self.mixed5b_pool_reduce = ReluLayer()
        self.mixed5b_3x3 = ReluLayer()
        self.mixed5b_5x5 = ReluLayer()
        self.mixed5b = CatLayer()
        self.softmax2 = SoftMaxLayer()

    def forward(self, x):
        conv2d0_pre_relu_conv_pad = F.pad(x, (2, 3, 2, 3))
        conv2d0_pre_relu_conv = self.conv2d0_pre_relu_conv(conv2d0_pre_relu_conv_pad)
        conv2d0         = self.conv2d0(conv2d0_pre_relu_conv)
        maxpool0_pad    = F.pad(conv2d0, (0, 1, 0, 1), value=float('-inf'))
        maxpool0        = self.maxpool0(maxpool0_pad, kernel_size=(3, 3), stride=(2, 2), padding=0, ceil_mode=False)
        localresponsenorm0 = F.local_response_norm(maxpool0, size=9, alpha=9.99999974738e-05, beta=0.5, k=1)
        conv2d1_pre_relu_conv = self.conv2d1_pre_relu_conv(localresponsenorm0)
        conv2d1         = self.conv2d1(conv2d1_pre_relu_conv)
        conv2d2_pre_relu_conv_pad = F.pad(conv2d1, (1, 1, 1, 1))
        conv2d2_pre_relu_conv = self.conv2d2_pre_relu_conv(conv2d2_pre_relu_conv_pad)
        conv2d2         = self.conv2d2(conv2d2_pre_relu_conv)
        localresponsenorm1 = F.local_response_norm(conv2d2, size=9, alpha=9.99999974738e-05, beta=0.5, k=1)
        maxpool1_pad    = F.pad(localresponsenorm1, (0, 1, 0, 1), value=float('-inf'))
        maxpool1        = self.maxpool1(maxpool1_pad, kernel_size=(3, 3), stride=(2, 2), padding=0, ceil_mode=False)
        mixed3a_1x1_pre_relu_conv = self.mixed3a_1x1_pre_relu_conv(maxpool1)
        mixed3a_3x3_bottleneck_pre_relu_conv = self.mixed3a_3x3_bottleneck_pre_relu_conv(maxpool1)
        mixed3a_5x5_bottleneck_pre_relu_conv = self.mixed3a_5x5_bottleneck_pre_relu_conv(maxpool1)
        mixed3a_pool_pad = F.pad(maxpool1, (1, 1, 1, 1), value=float('-inf'))
        mixed3a_pool    = self.mixed3a_pool(mixed3a_pool_pad, kernel_size=(3, 3), stride=(1, 1), padding=0, ceil_mode=False)
        mixed3a_1x1     = self.mixed3a_1x1(mixed3a_1x1_pre_relu_conv)
        mixed3a_3x3_bottleneck = self.mixed3a_3x3_bottleneck(mixed3a_3x3_bottleneck_pre_relu_conv)
        mixed3a_5x5_bottleneck = self.mixed3a_5x5_bottleneck(mixed3a_5x5_bottleneck_pre_relu_conv)
        mixed3a_pool_reduce_pre_relu_conv = self.mixed3a_pool_reduce_pre_relu_conv(mixed3a_pool)
        mixed3a_3x3_pre_relu_conv_pad = F.pad(mixed3a_3x3_bottleneck, (1, 1, 1, 1))
        mixed3a_3x3_pre_relu_conv = self.mixed3a_3x3_pre_relu_conv(mixed3a_3x3_pre_relu_conv_pad)
        mixed3a_5x5_pre_relu_conv_pad = F.pad(mixed3a_5x5_bottleneck, (2, 2, 2, 2))
        mixed3a_5x5_pre_relu_conv = self.mixed3a_5x5_pre_relu_conv(mixed3a_5x5_pre_relu_conv_pad)
        mixed3a_pool_reduce = self.mixed3a_pool_reduce(mixed3a_pool_reduce_pre_relu_conv)
        mixed3a_3x3     = self.mixed3a_3x3(mixed3a_3x3_pre_relu_conv)
        mixed3a_5x5     = self.mixed3a_5x5(mixed3a_5x5_pre_relu_conv)
        mixed3a         = self.mixed3a((mixed3a_1x1, mixed3a_3x3, mixed3a_5x5, mixed3a_pool_reduce), 1)
        mixed3b_1x1_pre_relu_conv = self.mixed3b_1x1_pre_relu_conv(mixed3a)
        mixed3b_3x3_bottleneck_pre_relu_conv = self.mixed3b_3x3_bottleneck_pre_relu_conv(mixed3a)
        mixed3b_5x5_bottleneck_pre_relu_conv = self.mixed3b_5x5_bottleneck_pre_relu_conv(mixed3a)
        mixed3b_pool_pad = F.pad(mixed3a, (1, 1, 1, 1), value=float('-inf'))
        mixed3b_pool    = self.mixed3b_pool(mixed3b_pool_pad, kernel_size=(3, 3), stride=(1, 1), padding=0, ceil_mode=False)
        mixed3b_1x1     = self.mixed3b_1x1(mixed3b_1x1_pre_relu_conv)
        mixed3b_3x3_bottleneck = self.mixed3b_3x3_bottleneck(mixed3b_3x3_bottleneck_pre_relu_conv)
        mixed3b_5x5_bottleneck = self.mixed3b_5x5_bottleneck(mixed3b_5x5_bottleneck_pre_relu_conv)
        mixed3b_pool_reduce_pre_relu_conv = self.mixed3b_pool_reduce_pre_relu_conv(mixed3b_pool)
        mixed3b_3x3_pre_relu_conv_pad = F.pad(mixed3b_3x3_bottleneck, (1, 1, 1, 1))
        mixed3b_3x3_pre_relu_conv = self.mixed3b_3x3_pre_relu_conv(mixed3b_3x3_pre_relu_conv_pad)
        mixed3b_5x5_pre_relu_conv_pad = F.pad(mixed3b_5x5_bottleneck, (2, 2, 2, 2))
        mixed3b_5x5_pre_relu_conv = self.mixed3b_5x5_pre_relu_conv(mixed3b_5x5_pre_relu_conv_pad)
        mixed3b_pool_reduce = self.mixed3b_pool_reduce(mixed3b_pool_reduce_pre_relu_conv)
        mixed3b_3x3     = self.mixed3b_3x3(mixed3b_3x3_pre_relu_conv)
        mixed3b_5x5     = self.mixed3b_5x5(mixed3b_5x5_pre_relu_conv)
        mixed3b         = self.mixed3b((mixed3b_1x1, mixed3b_3x3, mixed3b_5x5, mixed3b_pool_reduce), 1)
        maxpool4_pad    = F.pad(mixed3b, (0, 1, 0, 1), value=float('-inf'))
        maxpool4        = self.maxpool4(maxpool4_pad, kernel_size=(3, 3), stride=(2, 2), padding=0, ceil_mode=False)
        mixed4a_1x1_pre_relu_conv = self.mixed4a_1x1_pre_relu_conv(maxpool4)
        mixed4a_3x3_bottleneck_pre_relu_conv = self.mixed4a_3x3_bottleneck_pre_relu_conv(maxpool4)
        mixed4a_5x5_bottleneck_pre_relu_conv = self.mixed4a_5x5_bottleneck_pre_relu_conv(maxpool4)
        mixed4a_pool_pad = F.pad(maxpool4, (1, 1, 1, 1), value=float('-inf'))
        mixed4a_pool    = self.mixed4a_pool(mixed4a_pool_pad, kernel_size=(3, 3), stride=(1, 1), padding=0, ceil_mode=False)
        mixed4a_1x1     = self.mixed4a_1x1(mixed4a_1x1_pre_relu_conv)
        mixed4a_3x3_bottleneck = self.mixed4a_3x3_bottleneck(mixed4a_3x3_bottleneck_pre_relu_conv)
        mixed4a_5x5_bottleneck = self.mixed4a_5x5_bottleneck(mixed4a_5x5_bottleneck_pre_relu_conv)
        mixed4a_pool_reduce_pre_relu_conv = self.mixed4a_pool_reduce_pre_relu_conv(mixed4a_pool)
        mixed4a_3x3_pre_relu_conv_pad = F.pad(mixed4a_3x3_bottleneck, (1, 1, 1, 1))
        mixed4a_3x3_pre_relu_conv = self.mixed4a_3x3_pre_relu_conv(mixed4a_3x3_pre_relu_conv_pad)
        mixed4a_5x5_pre_relu_conv_pad = F.pad(mixed4a_5x5_bottleneck, (2, 2, 2, 2))
        mixed4a_5x5_pre_relu_conv = self.mixed4a_5x5_pre_relu_conv(mixed4a_5x5_pre_relu_conv_pad)
        mixed4a_pool_reduce = self.mixed4a_pool_reduce(mixed4a_pool_reduce_pre_relu_conv)
        mixed4a_3x3     = self.mixed4a_3x3(mixed4a_3x3_pre_relu_conv)
        mixed4a_5x5     = self.mixed4a_5x5(mixed4a_5x5_pre_relu_conv)
        mixed4a         = self.mixed4a((mixed4a_1x1, mixed4a_3x3, mixed4a_5x5, mixed4a_pool_reduce), 1)
        mixed4b_1x1_pre_relu_conv = self.mixed4b_1x1_pre_relu_conv(mixed4a)
        mixed4b_3x3_bottleneck_pre_relu_conv = self.mixed4b_3x3_bottleneck_pre_relu_conv(mixed4a)
        mixed4b_5x5_bottleneck_pre_relu_conv = self.mixed4b_5x5_bottleneck_pre_relu_conv(mixed4a)
        mixed4b_pool_pad = F.pad(mixed4a, (1, 1, 1, 1), value=float('-inf'))
        mixed4b_pool    = self.mixed4b_pool(mixed4b_pool_pad, kernel_size=(3, 3), stride=(1, 1), padding=0, ceil_mode=False)
        mixed4b_1x1     = self.mixed4b_1x1(mixed4b_1x1_pre_relu_conv)
        mixed4b_3x3_bottleneck = self.mixed4b_3x3_bottleneck(mixed4b_3x3_bottleneck_pre_relu_conv)
        mixed4b_5x5_bottleneck = self.mixed4b_5x5_bottleneck(mixed4b_5x5_bottleneck_pre_relu_conv)
        mixed4b_pool_reduce_pre_relu_conv = self.mixed4b_pool_reduce_pre_relu_conv(mixed4b_pool)
        mixed4b_3x3_pre_relu_conv_pad = F.pad(mixed4b_3x3_bottleneck, (1, 1, 1, 1))
        mixed4b_3x3_pre_relu_conv = self.mixed4b_3x3_pre_relu_conv(mixed4b_3x3_pre_relu_conv_pad)
        mixed4b_5x5_pre_relu_conv_pad = F.pad(mixed4b_5x5_bottleneck, (2, 2, 2, 2))
        mixed4b_5x5_pre_relu_conv = self.mixed4b_5x5_pre_relu_conv(mixed4b_5x5_pre_relu_conv_pad)
        mixed4b_pool_reduce = self.mixed4b_pool_reduce(mixed4b_pool_reduce_pre_relu_conv)
        mixed4b_3x3     = self.mixed4b_3x3(mixed4b_3x3_pre_relu_conv)
        mixed4b_5x5     = self.mixed4b_5x5(mixed4b_5x5_pre_relu_conv)
        mixed4b         = self.mixed4b((mixed4b_1x1, mixed4b_3x3, mixed4b_5x5, mixed4b_pool_reduce), 1)
        mixed4c_1x1_pre_relu_conv = self.mixed4c_1x1_pre_relu_conv(mixed4b)
        mixed4c_3x3_bottleneck_pre_relu_conv = self.mixed4c_3x3_bottleneck_pre_relu_conv(mixed4b)
        mixed4c_5x5_bottleneck_pre_relu_conv = self.mixed4c_5x5_bottleneck_pre_relu_conv(mixed4b)
        mixed4c_pool_pad = F.pad(mixed4b, (1, 1, 1, 1), value=float('-inf'))
        mixed4c_pool    = self.mixed4c_pool(mixed4c_pool_pad, kernel_size=(3, 3), stride=(1, 1), padding=0, ceil_mode=False)
        mixed4c_1x1     = self.mixed4c_1x1(mixed4c_1x1_pre_relu_conv)
        mixed4c_3x3_bottleneck = self.mixed4c_3x3_bottleneck(mixed4c_3x3_bottleneck_pre_relu_conv)
        mixed4c_5x5_bottleneck = self.mixed4c_5x5_bottleneck(mixed4c_5x5_bottleneck_pre_relu_conv)
        mixed4c_pool_reduce_pre_relu_conv = self.mixed4c_pool_reduce_pre_relu_conv(mixed4c_pool)
        mixed4c_3x3_pre_relu_conv_pad = F.pad(mixed4c_3x3_bottleneck, (1, 1, 1, 1))
        mixed4c_3x3_pre_relu_conv = self.mixed4c_3x3_pre_relu_conv(mixed4c_3x3_pre_relu_conv_pad)
        mixed4c_5x5_pre_relu_conv_pad = F.pad(mixed4c_5x5_bottleneck, (2, 2, 2, 2))
        mixed4c_5x5_pre_relu_conv = self.mixed4c_5x5_pre_relu_conv(mixed4c_5x5_pre_relu_conv_pad)
        mixed4c_pool_reduce = self.mixed4c_pool_reduce(mixed4c_pool_reduce_pre_relu_conv)
        mixed4c_3x3     = self.mixed4c_3x3(mixed4c_3x3_pre_relu_conv)
        mixed4c_5x5     = self.mixed4c_5x5(mixed4c_5x5_pre_relu_conv)
        mixed4c         = self.mixed4c((mixed4c_1x1, mixed4c_3x3, mixed4c_5x5, mixed4c_pool_reduce), 1)
        mixed4d_1x1_pre_relu_conv = self.mixed4d_1x1_pre_relu_conv(mixed4c)
        mixed4d_3x3_bottleneck_pre_relu_conv = self.mixed4d_3x3_bottleneck_pre_relu_conv(mixed4c)
        mixed4d_5x5_bottleneck_pre_relu_conv = self.mixed4d_5x5_bottleneck_pre_relu_conv(mixed4c)
        mixed4d_pool_pad = F.pad(mixed4c, (1, 1, 1, 1), value=float('-inf'))
        mixed4d_pool    = self.mixed4d_pool(mixed4d_pool_pad, kernel_size=(3, 3), stride=(1, 1), padding=0, ceil_mode=False)
        mixed4d_1x1     = self.mixed4d_1x1(mixed4d_1x1_pre_relu_conv)
        mixed4d_3x3_bottleneck = self.mixed4d_3x3_bottleneck(mixed4d_3x3_bottleneck_pre_relu_conv)
        mixed4d_5x5_bottleneck = self.mixed4d_5x5_bottleneck(mixed4d_5x5_bottleneck_pre_relu_conv)
        mixed4d_pool_reduce_pre_relu_conv = self.mixed4d_pool_reduce_pre_relu_conv(mixed4d_pool)
        mixed4d_3x3_pre_relu_conv_pad = F.pad(mixed4d_3x3_bottleneck, (1, 1, 1, 1))
        mixed4d_3x3_pre_relu_conv = self.mixed4d_3x3_pre_relu_conv(mixed4d_3x3_pre_relu_conv_pad)
        mixed4d_5x5_pre_relu_conv_pad = F.pad(mixed4d_5x5_bottleneck, (2, 2, 2, 2))
        mixed4d_5x5_pre_relu_conv = self.mixed4d_5x5_pre_relu_conv(mixed4d_5x5_pre_relu_conv_pad)
        mixed4d_pool_reduce = self.mixed4d_pool_reduce(mixed4d_pool_reduce_pre_relu_conv)
        mixed4d_3x3     = self.mixed4d_3x3(mixed4d_3x3_pre_relu_conv)
        mixed4d_5x5     = self.mixed4d_5x5(mixed4d_5x5_pre_relu_conv)
        mixed4d         = self.mixed4d((mixed4d_1x1, mixed4d_3x3, mixed4d_5x5, mixed4d_pool_reduce), 1)
        mixed4e_1x1_pre_relu_conv = self.mixed4e_1x1_pre_relu_conv(mixed4d)
        mixed4e_3x3_bottleneck_pre_relu_conv = self.mixed4e_3x3_bottleneck_pre_relu_conv(mixed4d)
        mixed4e_5x5_bottleneck_pre_relu_conv = self.mixed4e_5x5_bottleneck_pre_relu_conv(mixed4d)
        mixed4e_pool_pad = F.pad(mixed4d, (1, 1, 1, 1), value=float('-inf'))
        mixed4e_pool    = self.mixed4e_pool(mixed4e_pool_pad, kernel_size=(3, 3), stride=(1, 1), padding=0, ceil_mode=False)
        mixed4e_1x1     = self.mixed4e_1x1(mixed4e_1x1_pre_relu_conv)
        mixed4e_3x3_bottleneck = self.mixed4e_3x3_bottleneck(mixed4e_3x3_bottleneck_pre_relu_conv)
        mixed4e_5x5_bottleneck = self.mixed4e_5x5_bottleneck(mixed4e_5x5_bottleneck_pre_relu_conv)
        mixed4e_pool_reduce_pre_relu_conv = self.mixed4e_pool_reduce_pre_relu_conv(mixed4e_pool)
        mixed4e_3x3_pre_relu_conv_pad = F.pad(mixed4e_3x3_bottleneck, (1, 1, 1, 1))
        mixed4e_3x3_pre_relu_conv = self.mixed4e_3x3_pre_relu_conv(mixed4e_3x3_pre_relu_conv_pad)
        mixed4e_5x5_pre_relu_conv_pad = F.pad(mixed4e_5x5_bottleneck, (2, 2, 2, 2))
        mixed4e_5x5_pre_relu_conv = self.mixed4e_5x5_pre_relu_conv(mixed4e_5x5_pre_relu_conv_pad)
        mixed4e_pool_reduce = self.mixed4e_pool_reduce(mixed4e_pool_reduce_pre_relu_conv)
        mixed4e_3x3     = self.mixed4e_3x3(mixed4e_3x3_pre_relu_conv)
        mixed4e_5x5     = self.mixed4e_5x5(mixed4e_5x5_pre_relu_conv)
        mixed4e         = self.mixed4e((mixed4e_1x1, mixed4e_3x3, mixed4e_5x5, mixed4e_pool_reduce), 1)
        maxpool10_pad   = F.pad(mixed4e, (0, 1, 0, 1), value=float('-inf'))
        maxpool10       = self.maxpool10(maxpool10_pad, kernel_size=(3, 3), stride=(2, 2), padding=0, ceil_mode=False)
        mixed5a_1x1_pre_relu_conv = self.mixed5a_1x1_pre_relu_conv(maxpool10)
        mixed5a_3x3_bottleneck_pre_relu_conv = self.mixed5a_3x3_bottleneck_pre_relu_conv(maxpool10)
        mixed5a_5x5_bottleneck_pre_relu_conv = self.mixed5a_5x5_bottleneck_pre_relu_conv(maxpool10)
        mixed5a_pool_pad = F.pad(maxpool10, (1, 1, 1, 1), value=float('-inf'))
        mixed5a_pool    = self.mixed5a_pool(mixed5a_pool_pad, kernel_size=(3, 3), stride=(1, 1), padding=0, ceil_mode=False)
        mixed5a_1x1     = self.mixed5a_1x1(mixed5a_1x1_pre_relu_conv)
        mixed5a_3x3_bottleneck = self.mixed5a_3x3_bottleneck(mixed5a_3x3_bottleneck_pre_relu_conv)
        mixed5a_5x5_bottleneck = self.mixed5a_5x5_bottleneck(mixed5a_5x5_bottleneck_pre_relu_conv)
        mixed5a_pool_reduce_pre_relu_conv = self.mixed5a_pool_reduce_pre_relu_conv(mixed5a_pool)
        mixed5a_3x3_pre_relu_conv_pad = F.pad(mixed5a_3x3_bottleneck, (1, 1, 1, 1))
        mixed5a_3x3_pre_relu_conv = self.mixed5a_3x3_pre_relu_conv(mixed5a_3x3_pre_relu_conv_pad)
        mixed5a_5x5_pre_relu_conv_pad = F.pad(mixed5a_5x5_bottleneck, (2, 2, 2, 2))
        mixed5a_5x5_pre_relu_conv = self.mixed5a_5x5_pre_relu_conv(mixed5a_5x5_pre_relu_conv_pad)
        mixed5a_pool_reduce = self.mixed5a_pool_reduce(mixed5a_pool_reduce_pre_relu_conv)
        mixed5a_3x3     = self.mixed5a_3x3(mixed5a_3x3_pre_relu_conv)
        mixed5a_5x5     = self.mixed5a_5x5(mixed5a_5x5_pre_relu_conv)
        mixed5a         = self.mixed5a((mixed5a_1x1, mixed5a_3x3, mixed5a_5x5, mixed5a_pool_reduce), 1)
        mixed5b_1x1_pre_relu_conv = self.mixed5b_1x1_pre_relu_conv(mixed5a)
        mixed5b_3x3_bottleneck_pre_relu_conv = self.mixed5b_3x3_bottleneck_pre_relu_conv(mixed5a)
        mixed5b_5x5_bottleneck_pre_relu_conv = self.mixed5b_5x5_bottleneck_pre_relu_conv(mixed5a)
        mixed5b_pool_pad = F.pad(mixed5a, (1, 1, 1, 1), value=float('-inf'))
        mixed5b_pool    = self.mixed5b_pool(mixed5b_pool_pad, kernel_size=(3, 3), stride=(1, 1), padding=0, ceil_mode=False)
        mixed5b_1x1     = self.mixed5b_1x1(mixed5b_1x1_pre_relu_conv)
        mixed5b_3x3_bottleneck = self.mixed5b_3x3_bottleneck(mixed5b_3x3_bottleneck_pre_relu_conv)
        mixed5b_5x5_bottleneck = self.mixed5b_5x5_bottleneck(mixed5b_5x5_bottleneck_pre_relu_conv)
        mixed5b_pool_reduce_pre_relu_conv = self.mixed5b_pool_reduce_pre_relu_conv(mixed5b_pool)
        mixed5b_3x3_pre_relu_conv_pad = F.pad(mixed5b_3x3_bottleneck, (1, 1, 1, 1))
        mixed5b_3x3_pre_relu_conv = self.mixed5b_3x3_pre_relu_conv(mixed5b_3x3_pre_relu_conv_pad)
        mixed5b_5x5_pre_relu_conv_pad = F.pad(mixed5b_5x5_bottleneck, (2, 2, 2, 2))
        mixed5b_5x5_pre_relu_conv = self.mixed5b_5x5_pre_relu_conv(mixed5b_5x5_pre_relu_conv_pad)
        mixed5b_pool_reduce = self.mixed5b_pool_reduce(mixed5b_pool_reduce_pre_relu_conv)
        mixed5b_3x3     = self.mixed5b_3x3(mixed5b_3x3_pre_relu_conv)
        mixed5b_5x5     = self.mixed5b_5x5(mixed5b_5x5_pre_relu_conv)
        mixed5b         = self.mixed5b((mixed5b_1x1, mixed5b_3x3, mixed5b_5x5, mixed5b_pool_reduce), 1)
        avgpool0        = F.avg_pool2d(mixed5b, kernel_size=(7, 7), stride=(1, 1), padding=(0,), ceil_mode=False, count_include_pad=False)
        avgpool0_reshape = torch.reshape(input = avgpool0, shape = (-1,1024))
        softmax2_pre_activation_matmul = self.softmax2_pre_activation_matmul(avgpool0_reshape)
        softmax2        = self.softmax2(softmax2_pre_activation_matmul)
        return softmax2
	# tensorflow/lucid CPPN (X,Y) --> (R,G,B) Differentiable Image Parameterization in PyTorch
	import torch
	import torch.nn as nn
	import torch.optim as optim
	import torchvision.transforms as transforms
	from inception5h import Inception5h
	from PIL import Image



	def main():
	# Setup params
	image_size=(224,224)
	iterations = 512
	lr = 0.005
	use_device = 'cuda:0'

	# Uncomment to use seed
	#set_seed(876)

	# Target layer and channel
	layer_name='mixed4b_3x3_pre_relu_conv'
	channel=77

	# Load InceptionV1
	cnn = Inception5h()
	cnn.load_state_dict(torch.load('inception5h.pth'))
	cnn.add_layers() # Add layers not found in model file
	net = cnn.to(use_device).eval()
	for param in net.parameters():
	param.requires_grad = False

	# Uncomment to print all hookable layers
	#print(get_hookable_layers(cnn))

	# Create instance of CPPN
	img_cppn = CPPN_Conv(size=image_size, num_channels=24, num_layers=8, use_device=use_device, normalize=False)
	print('CPPN Params Count', sum(p.numel() for p in img_cppn.parameters() if p.requires_grad))

	loss_modules = register_hook_module(net, layer_name, channel)

	# Create 224x224 image
	output_tensor, img_cppn = dream_cppn(net, img_cppn, iterations, lr, loss_modules, use_device)
	simple_deprocess(output_tensor, name='out.png')

	# Create 720x720 image using multiscale generation
	image_size=(720,720)
	img_cppn.remake_input(image_size, use_device=use_device)
	output_tensor, img_cppn = dream_cppn(net, img_cppn, iterations, lr, loss_modules, use_device)
	simple_deprocess(output_tensor, name='out.png')



	# Simple Deprocess
	def simple_deprocess(output_tensor, name):
	output_tensor = output_tensor.squeeze(0).cpu() /255
	output_tensor.clamp_(0, 1)
	Image2PIL = transforms.ToPILImage()
	image = Image2PIL(output_tensor.cpu())
	image.save(name)


	# Function to maximize CNN activations
	def dream_cppn(net, img_cppn, iterations, lr, loss_modules, use_device):
	# Setup optimizer to optimize CPPN instance
	optimizer = torch.optim.Adam(img_cppn.parameters(), lr=lr)

	# Training loop
	for i in range(iterations):
	optimizer.zero_grad()
	img = img_cppn() * 255 # Create RGB image with CPPN
	net(img) # Create loss values
	loss = 0
	for mod in loss_modules: # Collect loss values
	loss += mod.loss
	loss.backward()

	# Uncomment to save iterations
	#if i % 25 == 0:
	# simple_deprocess(img.detach(), 'out_'+str(i)+'.png')

	print('Iteration', str(i+1), 'Loss', str(loss.item()))

	optimizer.step()
	img = img_cppn() * 255
	return img, img_cppn


	# Activation function for CPPN
	class CompositeActivation(nn.Module):
	def forward(self, input):
	input = torch.atan(input)
	return torch.cat([input / 0.67, (input * input) / 0.6], 1)


	# Compositional pattern-producing network (CPPN) with Conv2d layers
	class CPPN_Conv(nn.Module):

	def __init__(self, size=(224, 224), num_channels=24, num_layers=8, activ_func=CompositeActivation(), use_device='cpu', normalize=False):
	super(CPPN_Conv, self).__init__()
	self.input_size = size
	self.n_channels = num_channels
	self.net = self.create_net(num_channels, num_layers, activ_func, use_device, normalize)
	self.cppn_input = self.create_input(use_device)

	# Create CPPN (X,Y) --> (R,G,B)
	def create_net(self, num_channels, num_layers, activ_func, use_device, instance_norm, affine=True, bias=True):
	net = nn.Sequential()
	net.add_module(str(len(net)), nn.Conv2d(in_channels=2, out_channels=num_channels, kernel_size=1, bias=bias))
	if instance_norm:
	net.add_module(str(len(net)), nn.InstanceNorm2d(num_channels, affine=affine))
	net.add_module(str(len(net)), activ_func)
	for l in range(num_layers - 1):
	net.add_module(str(len(net)), nn.Conv2d(in_channels=num_channels*2, out_channels=num_channels, kernel_size=1, bias=bias))
	if instance_norm:
	net.add_module(str(len(net)), nn.InstanceNorm2d(num_channels, affine=affine))
	net.add_module(str(len(net)), activ_func)
	net.add_module(str(len(net)), nn.Conv2d(in_channels=num_channels*2, out_channels=3, kernel_size=1, bias=bias))
	net.add_module(str(len(net)), nn.Sigmoid())
	net.apply(self.cppn_normal)
	return net.to(use_device)

	# Create X,Y input for CPPN
	def create_input(self, use_device):
	if type(self.input_size) is not tuple and type(self.input_size) is not list:
	self.input_size = (self.input_size, self.input_size)
	w = torch.arange(0, self.input_size[1]).to(use_device)
	h = torch.arange(0, self.input_size[0]).to(use_device)
	w_exp = w.unsqueeze(1).expand((self.input_size[1], self.input_size[0])).true_divide(self.input_size[0]) - 0.5
	h_exp = h.unsqueeze(0).expand((self.input_size[1], self.input_size[0])).true_divide(self.input_size[1]) - 0.5
	return torch.stack((w_exp, h_exp), -1).permute(2,1,0).unsqueeze(0)

	# Replace input with different sized input
	def remake_input(self, image_size, use_device):
	self.input_size = image_size
	self.cppn_input = self.create_input(use_device)

	# Normalize Conv2d weights
	def cppn_normal(self, l):
	if type(l) == nn.Conv2d:
	l.weight.data.normal_(0, (1.0/self.n_channels)**(1/2))
	if l.bias is not None:
	nn.init.zeros_(l.bias)

	def forward(self):
	return self.net(self.cppn_input)


	# Create loss module and hook
	def register_hook_module(net, layer_name, channel=-1):
	loss_module = SimpleDreamLossHook(channel)
	getattr(net, layer_name).register_forward_hook(loss_module)
	return [loss_module]


	# Define a simple forward hook to collect DeepDream loss
	class SimpleDreamLossHook(nn.Module):
	def __init__(self, channel=-1):
	super(SimpleDreamLossHook, self).__init__()
	self.get_loss = self.mean_loss
	self.channel = channel

	def mean_loss(self, input):
	return input.mean()

	def forward(self, module, input, output):
	if self.channel > -1:
	self.loss = -self.get_loss(output[:,self.channel,:,:])
	else:
	self.loss = -self.get_loss(output)


	# Set global seeds to output reproducible
	def set_seed(seed):
	import random
	torch.manual_seed(seed)
	torch.cuda.manual_seed_all(seed)
	torch.backends.cudnn.deterministic=True
	random.seed(seed)


	# Get layers that can be hooked
	def get_hookable_layers(cnn):
	return [l[0] for l in list(cnn.named_children())]



	if __name__ == "__main__":
	main()