def __init__(self, data_size, latent_size=128, depth=3):
super().__init__()
upmode = 'bilinear'
c, h, w = data_size
cs = [c] + [2**(d+4) for d in range(depth)]
div = 2 ** depth
cl = lambda x : int(math.ceil(x))
modules = [
nn.Linear(latent_size, cs[-1] * cl(h/div) * cl(w/div)), nn.ReLU(),
util.Reshape( (cs[-1], cl(h/div), cl(w/div)) )
]
for d in range(depth, 0, -1):
modules += [
nn.Upsample(scale_factor=2, mode=upmode),
nn.ConvTranspose2d(cs[d], cs[d], 3, padding=1), nn.ReLU(),
nn.ConvTranspose2d(cs[d], cs[d-1], 3, padding=1), nn.ReLU()
]
modules += [
nn.ConvTranspose2d(c, c, (3, 3), padding=1), nn.Sigmoid(),
util.Lambda(lambda x : x[:, :, :h, :w]) # crop out any extra pixels due to rounding errors
]
self.decoder = nn.Sequential(*modules)
def __init__(self, insize, num_classes, mul=1.0, **kwargs):
super().__init__()
c, h, w = insize
self.prep = nn.Sequential(
conv((c, h, w),
64,
kernel_size=3,
padding=1,
bias=False,
sparse=False,
**kwargs), nn.BatchNorm2d(64), nn.ReLU())
self.layer0 = DavidLayer((64, h, w), 128, sparse=False,
**kwargs) # one maxpool
self.mid = nn.Sequential(
conv((128, h // 2, w // 2),
256,
kernel_size=3,
padding=1,
bias=False,
sparse=True,
**kwargs), nn.BatchNorm2d(256), nn.ReLU(),
nn.MaxPool2d(kernel_size=2))
self.layer1 = DavidLayer((256, h // 4, w // 4),
512,
sparse=False,
**kwargs) # one maxpool
self.head = nn.Sequential( # h//8, w//8
nn.MaxPool2d(kernel_size=4), util.Flatten(),
nn.Linear(512 * h // 32 * w // 32, num_classes),
util.Lambda(lambda x: x * mul))
def __init__(self, ni, nf, stride, drop_p=0.0, mult=0.2):
super().__init__()
self.mult = mult
self.bn = nn.BatchNorm2d(ni)
self.conv1 = nn.Conv2d(in_channels=ni,
out_channels=nf,
kernel_size=3,
stride=stride,
padding=1)
self.conv2 = bn_relu_conv(nf, nf, 3, 1)
self.drop = nn.Dropout(drop_p) if drop_p else None
self.shortcut = nn.Conv2d(
in_channels=ni, out_channels=nf, kernel_size=1,
stride=stride) if ni != nf else util.Lambda(lambda x: x)
def __init__(self,
data_size,
latent_size=(5, 5, 128),
depth=3,
gadditional=2,
radditional=4,
region=0.2,
method='clamp',
sigma_scale=1.0,
min_sigma=0.01):
super().__init__()
self.method, self.gadditional, self.radditional = method, gadditional, radditional
self.sigma_scale, self.min_sigma = sigma_scale, min_sigma
# latent space
self.latent = nn.Parameter(torch.randn(size=latent_size))
self.region = [int(r * region) for r in latent_size[:-1]]
ln = len(latent_size)
emb_size = latent_size[-1]
c, h, w = data_size
cs = [c] + [2**(d + 4) for d in range(depth)]
div = 2**depth
modules = []
for d in range(depth):
modules += [
nn.Conv2d(cs[d], cs[d + 1], 3, padding=1),
nn.ReLU(),
nn.Conv2d(cs[d + 1], cs[d + 1], 3, padding=1),
nn.ReLU(),
nn.MaxPool2d((2, 2))
]
modules += [
util.Flatten(),
nn.Linear(cs[-1] * (h // div) * (w // div), 1024),
nn.ReLU(),
nn.Linear(
1024, len(latent_size)
) # encoder produces a cont. index tuple (ln -1 for the means, 1 for the sigma)
]
self.encoder = nn.Sequential(*modules)
upmode = 'bilinear'
cl = lambda x: int(math.ceil(x))
modules = [
nn.Linear(emb_size, cs[-1] * cl(h / div) * cl(w / div)),
nn.ReLU(),
util.Reshape((cs[-1], cl(h / div), cl(w / div)))
]
for d in range(depth, 0, -1):
modules += [
nn.Upsample(scale_factor=2, mode=upmode),
nn.ConvTranspose2d(cs[d], cs[d], 3, padding=1),
nn.ReLU(),
nn.ConvTranspose2d(cs[d], cs[d - 1], 3, padding=1),
nn.ReLU()
]
modules += [
nn.ConvTranspose2d(c, c, (3, 3), padding=1),
nn.Sigmoid(),
util.Lambda(lambda x: x[:, :, :h, :w]
) # crop out any extra pixels due to rounding errors
]
self.decoder = nn.Sequential(*modules)
self.smp = True
def getmodel(arg, insize, numcls):
h1, h2 = arg.hidden
# if arg.method == 'l1' or arg.method == 'lp':
#
# one = nn.Linear(util.prod(insize), h1)
# two = nn.Linear(h1, h2)
# three = nn.Linear(h2, numcls)
#
# model = nn.Sequential(
# util.Flatten(),
# one, nn.Sigmoid(),
# two, nn.Sigmoid(),
# three, nn.Softmax()
# )
if arg.method == 'nas':
"""
Non-templated NAS model
"""
rng = getrng(arg.range[0], (h1, ) + insize)
c = arg.k[0]
one = NASLayer(in_size=insize,
out_size=(h1, ),
k=h1 * c,
gadditional=arg.gadditional[0],
radditional=arg.radditional[0],
region=rng,
has_bias=True,
fix_values=arg.fix_values,
min_sigma=arg.min_sigma,
template=None,
learn_cols=None,
chunk_size=c)
rng = getrng(arg.range[1], (h2, h1))
c = arg.k[1]
two = NASLayer(in_size=(h1, ),
out_size=(h2, ),
k=h2 * c,
gadditional=arg.gadditional[1],
radditional=arg.radditional[1],
region=rng,
has_bias=True,
fix_values=arg.fix_values,
min_sigma=arg.min_sigma,
template=None,
learn_cols=None,
chunk_size=c)
rng = getrng(arg.range[2], (numcls, h2))
c = arg.k[2]
three = NASLayer(in_size=(h2, ),
out_size=(numcls, ),
k=numcls * c,
gadditional=arg.gadditional[2],
radditional=arg.radditional[2],
region=rng,
has_bias=True,
fix_values=arg.fix_values,
min_sigma=arg.min_sigma,
template=None,
learn_cols=None,
chunk_size=c)
model = nn.Sequential(
one,
nn.Sigmoid(),
two,
nn.Sigmoid(),
three,
nn.Softmax(),
)
elif arg.method == 'nas-temp':
"""
Templated NAS model. Fixed output dimensions.
"""
rng = getrng(arg.range[0], (insize[1], insize[2]))
c = arg.k[0]
template = torch.arange(h1, dtype=torch.long)[:, None].expand(
h1, c).contiguous().view(h1 * c, 1)
template = torch.cat(
[template, torch.zeros(h1 * c, 3, dtype=torch.long)], dim=1)
one = NASLayer(in_size=insize,
out_size=(h1, ),
k=h1 * c,
gadditional=arg.gadditional[0],
radditional=arg.radditional[0],
region=rng,
has_bias=True,
fix_values=arg.fix_values,
min_sigma=arg.min_sigma,
template=template,
learn_cols=(1, 2, 3) if insize[0] > 1 else (2, 3),
chunk_size=c)
rng = getrng(arg.range[1], (h1, ))
c = arg.k[1]
template = torch.arange(h2, dtype=torch.long)[:, None].expand(
h2, c).contiguous().view(h2 * c, 1)
template = torch.cat(
[template, torch.zeros(h2 * c, 1, dtype=torch.long)], dim=1)
two = NASLayer(in_size=(h1, ),
out_size=(h2, ),
k=h2 * c,
gadditional=arg.gadditional[1],
radditional=arg.radditional[1],
region=rng,
has_bias=True,
fix_values=arg.fix_values,
min_sigma=arg.min_sigma,
template=template,
learn_cols=(1, ),
chunk_size=c)
rng = getrng(arg.range[2], (h2, ))
c = arg.k[2]
template = torch.arange(numcls, dtype=torch.long)[:, None].expand(
numcls, c).contiguous().view(numcls * c, 1)
template = torch.cat(
[template, torch.zeros(numcls * c, 1, dtype=torch.long)], dim=1)
three = NASLayer(in_size=(h2, ),
out_size=(numcls, ),
k=numcls * c,
gadditional=arg.gadditional[2],
radditional=arg.radditional[2],
region=rng,
has_bias=True,
fix_values=arg.fix_values,
min_sigma=arg.min_sigma,
template=template,
learn_cols=(1, ),
chunk_size=c)
model = nn.Sequential(
one,
nn.Sigmoid(),
two,
nn.Sigmoid(),
three,
nn.Softmax(),
)
elif arg.method == 'nas-conv':
"""
Convolutional NAS model.
"""
c1, c2 = h1, h2
one = Convolution(in_size=(1, 28, 28),
out_channels=c1,
k=arg.k[0],
kernel_size=7,
gadditional=arg.gadditional[0],
radditional=arg.radditional[1],
rprop=arg.range[0],
min_sigma=arg.min_sigma,
sigma_scale=arg.sigma_scale,
fix_values=arg.fix_values,
has_bias=True)
two = Convolution(in_size=(c1, 14, 14),
out_channels=c2,
k=arg.k[1],
kernel_size=7,
gadditional=arg.gadditional[1],
radditional=arg.radditional[1],
rprop=arg.range[1],
min_sigma=arg.min_sigma,
sigma_scale=arg.sigma_scale,
fix_values=arg.fix_values,
has_bias=True)
three = Convolution(in_size=(c2, 7, 7),
out_channels=numcls,
k=arg.k[2],
kernel_size=7,
gadditional=arg.gadditional[2],
radditional=arg.radditional[2],
rprop=arg.range[2],
min_sigma=arg.min_sigma,
sigma_scale=arg.sigma_scale,
fix_values=arg.fix_values,
has_bias=True)
model = nn.Sequential(
one,
nn.Sigmoid(),
nn.MaxPool2d(2),
two,
nn.Sigmoid(),
nn.MaxPool2d(2),
three,
nn.Sigmoid(),
nn.MaxPool2d(2),
util.Lambda(
lambda x: x.mean(dim=-1).mean(dim=-1)), # global average pool
nn.Softmax())
elif arg.method == 'conv':
c1, c2 = h1, h2
one = nn.Conv2d(insize[0], c1, kernel_size=3, padding=1)
two = nn.Conv2d(c1, c2, kernel_size=3, padding=1)
three = nn.Conv2d(c2, numcls, kernel_size=3, padding=1)
model = nn.Sequential(
one,
nn.Sigmoid(),
nn.MaxPool2d(2),
two,
nn.Sigmoid(),
nn.MaxPool2d(2),
three,
nn.Sigmoid(),
nn.MaxPool2d(2),
util.Lambda(
lambda x: x.mean(dim=-1).mean(dim=-1)), # global average pool
nn.Softmax())
elif arg.method == 'one':
"""
Convolutional NAS model.
"""
c1, c2 = h1, h2
one = Convolution(in_size=(1, 28, 28),
out_channels=c1,
k=arg.k[0],
kernel_size=7,
gadditional=arg.gadditional[0],
radditional=arg.radditional[1],
rprop=arg.range[0],
min_sigma=arg.min_sigma,
sigma_scale=arg.sigma_scale,
fix_values=arg.fix_values,
has_bias=True)
two = nn.Conv2d(c1, c2, kernel_size=3, padding=1)
three = nn.Conv2d(c2, numcls, kernel_size=3, padding=1)
model = nn.Sequential(
one,
nn.Sigmoid(),
nn.MaxPool2d(2),
two,
nn.Sigmoid(),
nn.MaxPool2d(2),
three,
nn.Sigmoid(),
nn.MaxPool2d(2),
util.Lambda(
lambda x: x.mean(dim=-1).mean(dim=-1)), # global average pool
nn.Softmax())
elif arg.method == 'two':
"""
Convolutional NAS model.
"""
c1, c2 = h1, h2
one = Convolution(in_size=(1, 28, 28),
out_channels=c1,
k=arg.k[0],
kernel_size=7,
gadditional=arg.gadditional[0],
radditional=arg.radditional[1],
rprop=arg.range[0],
min_sigma=arg.min_sigma,
sigma_scale=arg.sigma_scale,
fix_values=arg.fix_values,
has_bias=True)
two = Convolution(in_size=(c1, 14, 14),
out_channels=c2,
k=arg.k[1],
kernel_size=7,
gadditional=arg.gadditional[1],
radditional=arg.radditional[1],
rprop=arg.range[1],
min_sigma=arg.min_sigma,
sigma_scale=arg.sigma_scale,
fix_values=arg.fix_values,
has_bias=True)
three = nn.Conv2d(c2, numcls, kernel_size=3, padding=1)
model = nn.Sequential(
one,
nn.Sigmoid(),
nn.MaxPool2d(2),
two,
nn.Sigmoid(),
nn.MaxPool2d(2),
three,
nn.Sigmoid(),
nn.MaxPool2d(2),
util.Lambda(
lambda x: x.mean(dim=-1).mean(dim=-1)), # global average pool
nn.Softmax())
else:
raise Exception('Method {} not recognized'.format(arg.method))
if arg.cuda:
model.cuda()
return model, one, two, three