generator.py

import torch
import torch.nn as nn
import torchvision.models as models
import numpy as np

SEED = 3
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.cuda.manual_seed(SEED)


class GeneratorNet(nn.Module):
    def __init__(self, num_classes=80, norm=True, scale=True):
        super(GeneratorNet, self).__init__()
        self.add_info = AddInfo()
        self.generator = Generator()
        self.fc = nn.Linear(1024, num_classes, bias=False)
        self.s = nn.Parameter(torch.FloatTensor([10]))

    def forward(self, B1=None, B2=None, A=None, classifier=False):
        add_info = self.add_info(A, B1, B2)
        A_rebuild = self.generator(add_info)
        A_rebuild = self.l2_norm(A_rebuild)
        score = self.fc(A_rebuild*self.s)
        return A_rebuild, score
   
    def weight_norm(self):
        w = self.fc.weight.data
        norm = w.norm(p=2, dim=1, keepdim=True)
        self.fc.weight.data = w.div(norm.expand_as(w))

    def l2_norm(self,input):
        input_size = input.size()
        buffer = torch.pow(input, 2)

        norm = torch.sum(buffer, 1).add_(1e-10)
        norm = torch.sqrt(norm)

        _output = torch.div(input, norm.view(-1, 1).expand_as(input))

        output = _output.view(input_size)

        return output


class AddInfo(nn.Module):
    def __init__(self):
        super(AddInfo, self).__init__()
        self.fc = nn.Linear(1024, 2048)
        self.leakyRelu1 = nn.LeakyReLU(0.2, inplace=True)
        self.dropout = nn.Dropout(0.5)

    def forward(self, A, B1, B2):
        A = self.fc(A)
        A = self.leakyRelu1(A)
        B1 = self.fc(B1)
        B1 = self.leakyRelu1(B1)
        B2 = self.fc(B2)
        B2 = self.leakyRelu1(B2)
        out = A+(B1-B2)
        # print(torch.abs(B1-B2))
        out = self.dropout(out)

        return out

class Generator(nn.Module):
    def __init__(self):
        super(Generator, self).__init__()
        self.fc = nn.Linear(2048, 1024)
        self.leakyRelu = nn.LeakyReLU(0.2, inplace=True)
        self.dropout = nn.Dropout(0.5)

    def forward(self, x):
        out = self.fc(x)
        out = self.leakyRelu(out)
        out = self.dropout(out)

        return out


if __name__ == '__main__':
    A = torch.ones((64,1024))
    B1 = torch.ones((64,1024))
    B2 = torch.ones((64,1024))
    gen_net = GeneratorNet()
    gen_data, score = gen_net(A, B1, B2)
    print(gen_data.shape)
    print(score.shape)