Spaces:
Sleeping
Sleeping
File size: 2,657 Bytes
165ee00 |
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 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
# Author: Qiuyi
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils.parametrizations import spectral_norm
class _Combo(nn.Module):
def forward(self, input):
return self.model(input)
class LinearCombo(_Combo):
def __init__(self, in_features, out_features, activation=nn.LeakyReLU(0.2)):
super().__init__()
self.model = nn.Sequential(
nn.Linear(in_features, out_features),
activation
)
class SNLinearCombo(_Combo):
def __init__(self, in_features, out_features, activation=nn.LeakyReLU(0.2)):
super().__init__()
self.model = nn.Sequential(
spectral_norm(nn.Linear(in_features, out_features)),
activation
)
class MLP(nn.Module):
"""Regular fully connected network generating features.
Args:
in_features: The number of input features.
out_feature: The number of output features.
layer_width: The widths of the hidden layers.
combo: The layer combination to be stacked up.
Shape:
- Input: `(N, H_in)` where H_in = in_features.
- Output: `(N, H_out)` where H_out = out_features.
"""
def __init__(
self, in_features: int, out_features:int, layer_width: list,
combo = LinearCombo
):
super().__init__()
self.in_features = in_features
self.out_features = out_features
self.layer_width = list(layer_width)
self.model = self._build_model(combo)
def forward(self, input):
return self.model(input)
def _build_model(self, combo):
model = nn.Sequential()
idx = -1
for idx, (in_ftr, out_ftr) in enumerate(self.layer_sizes[:-1]):
model.add_module(str(idx), combo(in_ftr, out_ftr))
model.add_module(str(idx+1), nn.Linear(*self.layer_sizes[-1])) # type:ignore
return model
@property
def layer_sizes(self):
return list(zip([self.in_features] + self.layer_width,
self.layer_width + [self.out_features]))
class SNMLP(MLP):
def __init__(
self, in_features: int, out_features: int, layer_width: list,
combo=SNLinearCombo):
super().__init__(in_features, out_features, layer_width, combo)
def _build_model(self, combo):
model = nn.Sequential()
idx = -1
for idx, (in_ftr, out_ftr) in enumerate(self.layer_sizes[:-1]):
model.add_module(str(idx), combo(in_ftr, out_ftr))
model.add_module(str(idx+1), spectral_norm(nn.Linear(*self.layer_sizes[-1]))) # type:ignore
return model |