<commit_before>from __future__ import division
from __future__ import print_function
from __future__ import absolute_import

import os
import shutil

import torch
import torch.nn as nn
import torch.nn.functional as F


class SketchNet(nn.Module):
    """Network that takes as input sketch and photo vectors and computes 
    and adapted & normalized Euclidean distance between each photo and sketch.

    @param n_photos: integer [default: 32]
    """
    def __init__(self, layer='fc6', n_photos=32):
        super(SketchNet, self).__init__()
        if layer == 'conv42':
            self.sketch_adaptor = Conv42AdaptorNet()
            self.photo_adaptor = Conv42AdaptorNet()
        elif layer == 'fc6':
            self.sketch_adaptor = FC6AdaptorNet()
            self.photo_adaptor = FC6AdaptorNet()
        else:
            raise Exception('%s layer not supported.' % layer)
        # self.combiner = nn.Linear((n_photos + 1) * 100, 32)
        self.n_photos = n_photos
        self.layer = layer

    def forward(self, photos, sketch):
        batch_size = photos.size(0)
        assert photos.size(1) == self.n_photos
        # photos is a torch.Tensor of size batch_size x 32 x 4096 
        # sketch is a torch.Tensor of size batch_size x 4096 (single sketch)
        sketch = self.sketch_adaptor(sketch)
        photos = torch.cat([self.photo_adaptor(photos[:, i]).unsqueeze(1)
                            for i in xrange(self.n_photos)], dim=1)
        # compute euclidean distance from sketch to each photo
        # distances = torch.cat([torch.norm(photos[:, i] - sketch, p=2, dim=1).unsqueeze(1) 
        #                        for i in xrange(self.n_photos)], dim=1)
        distances = torch.cat([pearson_correlation(photos[:, i], sketch, dim=1).unsqueeze(1)
                               for i in xrange(self.n_photos)], dim=1)
        # distances = torch.cat((photos, sketch.unsqueeze(1)), dim=1)
        # distances = distances.view(batch_size, -1)
        # distances = self.combiner(distances)
        log_distances = F.log_softmax(distances, dim=1)
        return log_distances


class SketchNetHARD(SketchNet):
    def forward(self, photos, sketch):
        batch_size = photos.size(0)
        assert photos.size(1) == self.n_photos
        # photos is a torch.Tensor of size batch_size x 32 x 4096
        # sketch is a torch.Tensor of size batch_size x 4096 (single sketch)
        sketch = self.sketch_adaptor(sketch)
        photos = torch.cat([self.photo_adaptor(photos[:, i]).unsqueeze(1)
                            for i in xrange(self.n_photos)], dim=1)
        distances = torch.cat([pearson_correlation(photos[:, i], sketch, dim=1).unsqueeze(1)
                               for i in xrange(self.n_photos)], dim=1)
        return distances


class SketchNetCATEGORY(SketchNet):
    def __init__(self, layer='fc6', n_photos=32):
        super(SketchNetCATEGORY, self).__init__(layer=layer, n_photos=n_photos)
        self.annotation_net = AnnotationNet()
    
    # this only works for SINGLE photo and SINGLE sketch
    def forward(self, photo, sketch):
        batch_size = photo.size(0)
        sketch = self.sketch_adaptor(sketch)
        photo = self.photo_adaptor(photo)
        output = pearson_correlation(photo, sketch)
        annotation = self.annotation_net(sketch[:int(batch_size / 4)])  # only sketch
        return F.sigmoid(output), F.log_softmax(annotation)


class Conv42AdaptorNet(nn.Module):
    """Light network to learn an adapted embedding."""
    def __init__(self):
        super(Conv42AdaptorNet, self).__init__()
        self.conv_layers = nn.Sequential(
            nn.Conv2d(512, 16, kernel_size=5, stride=1),
            nn.LeakyReLU(),
            nn.MaxPool2d(2, stride=2, dilation=1))
        self.fc_layers = nn.Sequential(
            nn.Linear(16 * 12 * 12, 1000),
            nn.LeakyReLU(),
            nn.Dropout(),
            nn.Linear(1000, 1000))

    def forward(self, x):
        x = self.conv_layers(x)
        x = x.view(-1, 16 * 12 * 12)<commit_msg>I can reproduce the results now; and lets try it with new data<commit_after>from __future__ import division
from __future__ import print_function
from __future__ import absolute_import

import os
import shutil

import torch
import torch.nn as nn
import torch.nn.functional as F


class SketchNet(nn.Module):
    """Network that takes as input sketch and photo vectors and computes 
    and adapted & normalized Euclidean distance between each photo and sketch.

    @param n_photos: integer [default: 32]
    """
    def __init__(self, layer='fc6', n_photos=32):
        super(SketchNet, self).__init__()
        if layer == 'conv42':
            self.sketch_adaptor = Conv42AdaptorNet()
            self.photo_adaptor = Conv42AdaptorNet()
        elif layer == 'fc6':
            self.sketch_adaptor = FC6AdaptorNet()
            self.photo_adaptor = FC6AdaptorNet()
        else:
            raise Exception('%s layer not supported.' % layer)
        # self.combiner = nn.Linear((n_photos + 1) * 100, 32)
        self.n_photos = n_photos
        self.layer = layer

    def forward(self, photos, sketch):
        batch_size = photos.size(0)
        assert photos.size(1) == self.n_photos
        # photos is a torch.Tensor of size batch_size x 32 x 4096 
        # sketch is a torch.Tensor of size batch_size x 4096 (single sketch)
        sketch = self.sketch_adaptor(sketch)
        photos = torch.cat([self.photo_adaptor(photos[:, i]).unsqueeze(1)
                            for i in xrange(self.n_photos)], dim=1)
        # compute euclidean distance from sketch to each photo
        # distances = torch.cat([torch.norm(photos[:, i] - sketch, p=2, dim=1).unsqueeze(1) 
        #                        for i in xrange(self.n_photos)], dim=1)
        distances = torch.cat([pearson_correlation(photos[:, i], sketch, dim=1).unsqueeze(1)
                               for i in xrange(self.n_photos)], dim=1)
        # distances = torch.cat((photos, sketch.unsqueeze(1)), dim=1)
        # distances = distances.view(batch_size, -1)
        # distances = self.combiner(distances)
        log_distances = F.log_softmax(distances, dim=1)
        return log_distances


class SketchNetHARD(SketchNet):
    def forward(self, photos, sketch):
        batch_size = photos.size(0)
        assert photos.size(1) == self.n_photos
        # photos is a torch.Tensor of size batch_size x 32 x 4096
        # sketch is a torch.Tensor of size batch_size x 4096 (single sketch)
        sketch = self.sketch_adaptor(sketch)
        photos = torch.cat([self.photo_adaptor(photos[:, i]).unsqueeze(1)
                            for i in xrange(self.n_photos)], dim=1)
        distances = torch.cat([pearson_correlation(photos[:, i], sketch, dim=1).unsqueeze(1)
                               for i in xrange(self.n_photos)], dim=1)
        return distances


class SketchNetCATEGORY(SketchNet):
    def __init__(self, layer='fc6', n_photos=32):
        super(SketchNetCATEGORY, self).__init__(layer=layer, n_photos=n_photos)
        self.annotation_net = AnnotationNet()
    
    # this only works for SINGLE photo and SINGLE sketch
    def forward(self, photo, sketch):
        batch_size = photo.size(0)
        sketch = self.sketch_adaptor(sketch)
        photo = self.photo_adaptor(photo)
        output = pearson_correlation(photo, sketch)
        annotation = self.annotation_net(sketch[:int(batch_size / 4)])  # only sketch
        return F.sigmoid(output), F.log_softmax(annotation, dim=1)


class Conv42AdaptorNet(nn.Module):
    """Light network to learn an adapted embedding."""
    def __init__(self):
        super(Conv42AdaptorNet, self).__init__()
        self.conv_layers = nn.Sequential(
            nn.Conv2d(512, 64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            Swish(),
            nn.MaxPool2d(2, stride=2, dilation=1))
        self.fc_layers = nn.Sequential(
            nn.Linear(64 * 14 * 14, 4096),
            nn.BatchNorm1d(4096),
            Swish(),
            nn.Dropout(.5),
            nn.Linear(4096, 1000))

    def forward(self, x):
        x = self.conv_layers(x)
        x = x.view(-1, 64 * 14 * 14)
{'fraction_non_alphanumeric': 0.10830589096231248, 'fraction_numerical': 0.027686303207708256, 'mean_word_length': 2.4982935153583616, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3300644', 'n_tokens_mistral': 2751, 'n_tokens_neox': 2523, 'n_words': 641}