File size: 1,218 Bytes
02c5426 |
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 |
import math
from argparse import Namespace
import torch
import torch.nn as nn
import torch.nn.functional as F
import models
from models import register
import numpy as np
class ExpansionNet(nn.Module):
def __init__(self, args):
super(ExpansionNet, self).__init__()
self.args = args
self.in_dim = args.in_dim
self.out_dim = args.out_dim
self.hidden_list = args.hidden_list
layers = []
lastv = self.in_dim
hidden_list = self.hidden_list
out_dim = self.out_dim
for hidden in hidden_list:
layers.append(nn.Linear(lastv, hidden))
layers.append(nn.ReLU())
lastv = hidden
layers.append(nn.Linear(lastv, out_dim))
self.layers = nn.Sequential(*layers)
def forward(self, x):
b, _, c = x.shape
x = x.view(-1, c)
logits = self.layers(x)
out = nn.functional.normalize(logits, dim=1)
return out.view(b,_,self.out_dim)
@register('ExpansionNet')
def make_ExpansionNet(in_dim=580,out_dim=10,hidden_list=None):
args = Namespace()
args.in_dim = in_dim
args.out_dim = out_dim
args.hidden_list = hidden_list
return ExpansionNet(args)
|