def __init__(self, args, downsample):
super(IAFLayer, self).__init__()
n_in = args.h_size
n_out = args.h_size * 2 + args.z_size * 2
self.z_size = args.z_size
self.h_size = args.h_size
self.iaf = args.iaf
self.ds = downsample
self.args = args
if downsample:
stride, padding, filter_size = 2, 1, 4
self.down_conv_b = wn(
nn.ConvTranspose2d(args.h_size + args.z_size, args.h_size, 4,
2, 1))
else:
stride, padding, filter_size = 1, 1, 3
self.down_conv_b = wn(
nn.Conv2d(args.h_size + args.z_size, args.h_size, 3, 1, 1))
# create modules for UP pass:
self.up_conv_a = wn(
nn.Conv2d(n_in, n_out, filter_size, stride, padding))
self.up_conv_b = wn(nn.Conv2d(args.h_size, args.h_size, 3, 1, 1))
# create modules for DOWN pass:
self.down_conv_a = wn(
nn.Conv2d(n_in, 4 * self.z_size + 2 * self.h_size, 3, 1, 1))
if args.iaf:
self.down_ar_conv = ARMultiConv2d([args.h_size] * 2,
[args.z_size] * 2, args)
def __init__(self,
input_dim,
output_dim,
dk,
dv,
num_heads,
kernel_size,
padding,
rel_encoding=True,
height=None,
width=None,
norm='weight_norm'):
super(AttentionConv2d, self).__init__()
self.input_dim = input_dim
self.output_dim = output_dim
self.dk = dk
self.dv = dv
self.num_heads = num_heads
self.kernel_size = kernel_size
self.dkh = self.dk // self.num_heads
if rel_encoding and not height:
raise (
"Cannot use relative encoding without specifying input's height and width"
)
self.H = height
self.W = width
self.conv_qkv = nn.Conv2d(input_dim, 2 * dk + dv, 1)
self.conv_attn = nn.Conv2d(dv, dv, 1)
self.conv_out = nn.Conv2d(input_dim,
output_dim - dv,
kernel_size,
padding=padding)
self.softmax = nn.Softmax(dim=-1)
if rel_encoding:
self.key_rel_w = nn.Parameter(self.dkh**-0.5 +
torch.rand(2 * width - 1, self.dkh),
requires_grad=True)
self.key_rel_h = nn.Parameter(self.dkh**-0.5 +
torch.rand(2 * height - 1, self.dkh),
requires_grad=True)
self.relative_encoding = rel_encoding
if norm == 'weight_norm':
self.conv_qkv = wn(self.conv_qkv)
self.conv_attn = wn(self.conv_attn)
self.conv_out = wn(self.conv_out)
def __init__(self, num_filters_in, num_filters_out, filter_size=(2, 2), stride=(1, 1),
shift_output_right=False, augment=False, norm='weight_norm'):
super(down_right_shifted_conv2d, self).__init__()
assert norm in [None, 'batch_norm', 'weight_norm']
self.pad = nn.ZeroPad2d((filter_size[1] - 1, 0, filter_size[0] - 1, 0))
# self.augmented = augment
# if self.augmented:
# self.conv = AttentionConv2d(num_filters_in, num_filters_out, kernel_size=filter_size,
# dk=64, dv=64, num_heads=2, padding=0,
# height=32, width=32, norm=norm)
# else:
# self.conv = nn.Conv2d(num_filters_in, num_filters_out, filter_size, stride)
self.conv = nn.Conv2d(num_filters_in, num_filters_out, filter_size, stride)
self.shift_output_right = shift_output_right
self.norm = norm
if norm == 'weight_norm':
self.conv = wn(self.conv)
elif norm == 'batch_norm':
self.bn = nn.BatchNorm2d(num_filters_out)
if shift_output_right:
self.right_shift = lambda x: right_shift(x, pad=nn.ZeroPad2d((1, 0, 0, 0)))
def __init__(self,
num_filters_in,
num_filters_out,
filter_size=(2, 2),
stride=(1, 1),
shift_output_right=False,
norm='weight_norm'):
super(down_right_shifted_conv2d, self).__init__()
assert norm in [None, 'batch_norm', 'weight_norm']
self.pad = nn.ZeroPad2d((filter_size[1] - 1, 0, filter_size[0] - 1, 0))
self.conv = nn.Conv2d(num_filters_in,
num_filters_out,
filter_size,
stride=stride)
self.shift_output_right = shift_output_right
self.norm = norm
if norm == 'weight_norm':
self.conv = wn(self.conv)
elif norm == 'batch_norm':
self.bn = nn.BatchNorm2d(num_filters_out)
if shift_output_right:
self.right_shift = lambda x: right_shift(x,
pad=nn.ZeroPad2d(
(1, 0, 0, 0)))
def __init__(self, dim_in, dim_out):
super(nin, self).__init__()
self.lin_a = wn(nn.Linear(dim_in, dim_out))
self.dim_out = dim_out
self.conv2 = nn.Conv2d(in_channels=dim_in,
out_channels=dim_out,
kernel_size=1)
def __init__(self,
num_filters_in,
num_filters_out,
filter_size=(2, 3),
stride=(1, 1),
shift_output_down=False,
norm='weight_norm'):
super(down_shifted_conv2d, self).__init__(
) # Nawid - Padding seems to agree with TF version of official implementation
assert norm in [None, 'batch_norm', 'weight_norm']
self.conv = nn.Conv2d(num_filters_in, num_filters_out, filter_size,
stride) # Nawid - Convolution object
self.shift_output_down = shift_output_down
self.norm = norm
self.pad = nn.ZeroPad2d((
int((filter_size[1] - 1) / 2), # pad left
int((filter_size[1] - 1) / 2), # pad right
filter_size[0] - 1, # pad top
0)) # pad down
if norm == 'weight_norm':
self.conv == wn(
self.conv) # Nawid- Normalise the convolution object
elif norm == 'batch_norm':
self.bn = nn.BatchNorm2d(num_filters_out)
if shift_output_down: # Nawid - If shift output down, then the x is shifted downwards which is performed by down_shift method in the utils and x is given certain padding
self.down_shift = lambda x: down_shift(x,
pad=nn.ZeroPad2d(
(0, 0, 1, 0)))
def __init__(self,
num_filters_in,
num_filters_out,
filter_size=(2, 2),
stride=(1, 1),
shift_output_right=False,
norm='weight_norm'):
super(down_right_shifted_conv2d, self).__init__()
assert norm in [None, 'batch_norm', 'weight_norm']
self.pad = nn.ZeroPad2d(
(filter_size[1] - 1, 0, filter_size[0] - 1, 0)
) # Nawid -Adds Padding to the left and the top ( no padding to the bottom)
self.conv = nn.Conv2d(num_filters_in,
num_filters_out,
filter_size,
stride=stride) # Nawid - Convolution object
self.shift_output_right = shift_output_right # Nawid - Whether the output should be shifted
self.norm = norm
if norm == 'weight_norm':
self.conv == wn(
self.conv) #Nawid - Weighnormalisation on the convolution
elif norm == 'batch_norm':
self.bn = nn.BatchNorm2d(num_filters_out)
if shift_output_right:
self.right_shift = lambda x: right_shift(
x, pad=nn.ZeroPad2d(
(1, 0, 0, 0))) # Nawid - Performs right shift and padding
def __init__(self, num_filters_in, num_filters_out, filter_size=(2,2), stride=(1,1),
shift_output_right=False):
super(down_right_shifted_deconv2d, self).__init__()
self.deconv = wn(nn.ConvTranspose2d(num_filters_in, num_filters_out, filter_size,
stride, output_padding=1))
self.filter_size = filter_size
self.stride = stride
def __init__(self, num_filters_in, num_filters_out, filter_size=(2,2), stride=(1,1),
shift_output_right=False):
super(down_right_shifted_deconv2d, self).__init__()
self.deconv = wn(nn.ConvTranspose2d(num_filters_in, num_filters_out, filter_size,
stride, output_padding=1))
self.filter_size = filter_size
self.stride = stride
def __init__(self,
num_filters_in,
num_filters_out,
filter_size=(2, 3),
stride=(1, 1),
shift_output_down=False,
norm='weight_norm'):
super(down_shifted_conv2d, self).__init__()
assert norm in [None, 'batch_norm', 'weight_norm']
self.conv = nn.Conv2d(num_filters_in, num_filters_out, filter_size,
stride)
self.shift_output_down = shift_output_down
self.norm = norm
self.pad = nn.ZeroPad2d((
int((filter_size[1] - 1) / 2), # pad left
int((filter_size[1] - 1) / 2), # pad right
filter_size[0] - 1, # pad top
0)) # pad down
if norm == 'weight_norm':
self.conv = wn(self.conv)
elif norm == 'batch_norm':
self.bn = nn.BatchNorm2d(num_filters_out)
if shift_output_down:
self.down_shift = lambda x: down_shift(x,
pad=nn.ZeroPad2d(
(0, 0, 1, 0)))
def __init__(self, dim_in, dim_out, weight_norm=True):
super(nin, self).__init__()
if weight_norm:
self.lin_a = wn(nn.Linear(dim_in, dim_out))
else:
self.lin_a = nn.Linear(dim_in, dim_out)
self.dim_out = dim_out
def __init__(self, dim_in, dim_out, attention=False, height=32, width=32):
super(nin, self).__init__()
self.lin_a = wn(nn.Linear(dim_in, dim_out))
self.dim_out = dim_out
self.add_attention = attention
if attention:
self.pos = PositionwiseFeedForward(height * width, height * width)
self.attn = MultiHeadAttention(4, height * width, 64, 64)
def __init__(self,
num_filters_in,
num_filters_out,
filter_size=(2, 3),
stride=(1, 1)):
super(down_shifted_deconv2d, self).__init__()
self.deconv = wn(
nn.ConvTranspose2d(num_filters_in,
num_filters_out,
filter_size,
stride,
output_padding=1)
) # Nawid - Deconvolution object which is weight normalised
self.filter_size = filter_size
self.stride = stride
def __init__(self,
input_size,
num_classes,
dense_neurons,
weight_init=True):
super(Generator, self).__init__()
self.logger = logging.getLogger(__name__) # initialize logger
self.num_classes = num_classes
self.Dense = nn.Linear(input_size, dense_neurons)
self.Relu = nn.ReLU()
self.Tanh = nn.Tanh()
self.Deconv2D_0 = nn.ConvTranspose2d(in_channels=522,
out_channels=256,
kernel_size=5,
stride=2,
padding=2,
output_padding=1,
bias=False)
self.Deconv2D_1 = nn.ConvTranspose2d(in_channels=266,
out_channels=128,
kernel_size=5,
stride=2,
padding=2,
output_padding=1,
bias=False)
self.Deconv2D_2 = wn(
nn.ConvTranspose2d(in_channels=138,
out_channels=3,
kernel_size=5,
stride=2,
padding=2,
output_padding=1,
bias=False))
self.BatchNorm1D = nn.BatchNorm1d(dense_neurons)
self.BatchNorm2D_0 = nn.BatchNorm2d(256)
self.BatchNorm2D_1 = nn.BatchNorm2d(128)
if weight_init:
# initialize weights for all conv and lin layers
self.apply(init_weights)
# log network structure
self.logger.debug(self)
def __init__(self, num_filters_in, num_filters_out, filter_size=(2,2), stride=(1,1),
shift_output_right=False, norm='weight_norm'):
super(down_right_shifted_conv2d, self).__init__()
assert norm in [None, 'batch_norm', 'weight_norm']
self.pad = nn.ZeroPad2d((filter_size[1] - 1, 0, filter_size[0] - 1, 0))
self.conv = nn.Conv2d(num_filters_in, num_filters_out, filter_size, stride=stride)
self.shift_output_right = shift_output_right
self.norm = norm
if norm == 'weight_norm':
self.conv == wn(self.conv)
elif norm == 'batch_norm':
self.bn = nn.BatchNorm2d(num_filters_out)
if shift_output_right :
self.right_shift = lambda x : right_shift(x, pad=nn.ZeroPad2d((1, 0, 0, 0)))
def __init__(self,
num_filters_in,
num_filters_out,
filter_size=(2, 3),
stride=(1, 1),
shift_output_down=False,
augment=False,
norm='weight_norm'):
super(down_shifted_conv2d, self).__init__()
assert norm in [None, 'batch_norm', 'weight_norm']
self.augmented = augment
if self.augmented:
self.conv = AttentionConv2d(num_filters_in,
num_filters_out,
kernel_size=filter_size,
dk=64,
dv=64,
num_heads=2,
padding=0,
height=32,
width=32,
norm=norm)
else:
self.conv = nn.Conv2d(num_filters_in, num_filters_out, filter_size,
stride)
self.shift_output_down = shift_output_down
self.norm = norm
self.pad = nn.ZeroPad2d((
int((filter_size[1] - 1) / 2), # pad left
int((filter_size[1] - 1) / 2), # pad right
filter_size[0] - 1, # pad top
0)) # pad down
if norm == 'weight_norm' and not self.augmented:
self.conv = wn(self.conv)
elif norm == 'batch_norm':
self.bn = nn.BatchNorm2d(num_filters_out)
if shift_output_down:
self.down_shift = lambda x: down_shift(x,
pad=nn.ZeroPad2d(
(0, 0, 1, 0)))
def __init__(self, num_filters_in, num_filters_out, filter_size=(2,3), stride=(1,1),
shift_output_down=False, norm='weight_norm'):
super(down_shifted_conv2d, self).__init__()
assert norm in [None, 'batch_norm', 'weight_norm']
self.conv = nn.Conv2d(num_filters_in, num_filters_out, filter_size, stride)
self.shift_output_down = shift_output_down
self.norm = norm
self.pad = nn.ZeroPad2d((int((filter_size[1] - 1) / 2), # pad left
int((filter_size[1] - 1) / 2), # pad right
filter_size[0] - 1, # pad top
0) ) # pad down
if norm == 'weight_norm':
self.conv == wn(self.conv)
elif norm == 'batch_norm':
self.bn = nn.BatchNorm2d(num_filters_out)
if shift_output_down :
self.down_shift = lambda x : down_shift(x, pad=nn.ZeroPad2d((0, 0, 1, 0)))
def __init__(self,
num_filters_in,
num_filters_out,
filter_size=(2, 3),
stride=(1, 1),
norm=None,
num_actions=0):
super(down_shifted_deconv2d, self).__init__()
self.deconv = nn.ConvTranspose2d(num_filters_in + num_actions,
num_filters_out,
filter_size,
stride,
output_padding=1)
if norm == 'weight_norm':
self.deconv = wn(
nn.ConvTranspose2d(num_filters_in + num_actions,
num_filters_out,
filter_size,
stride,
output_padding=1))
self.filter_size = filter_size
self.stride = stride
def __init__(self, num_filters_in, num_filters_out, filter_size=(2, 3), stride=(1, 1),
shift_output_down=False, norm='weight_norm', attention=False):
super(down_shifted_conv2d, self).__init__()
assert norm in [None, 'batch_norm', 'weight_norm']
self.add_attention = attention
if attention:
h = 32
w = 32
self.pos = PositionwiseFeedForward(h * w, h * w)
self.attn = MultiHeadAttention(4, h * w, 64, 64)
# self.augmented = augment
# if self.augmented:
# self.conv = AttentionConv2d(num_filters_in, num_filters_out, kernel_size=filter_size,
# dk=64, dv=64, num_heads=2, padding=0, rel_encoding=False,
# height=32, width=32, norm=norm)
# else:
# self.conv = nn.Conv2d(num_filters_in, num_filters_out, filter_size, stride)
self.conv = nn.Conv2d(num_filters_in, num_filters_out, filter_size, stride)
self.shift_output_down = shift_output_down
self.norm = norm
self.pad = nn.ZeroPad2d((int((filter_size[1] - 1) / 2), # pad left
int((filter_size[1] - 1) / 2), # pad right
filter_size[0] - 1, # pad top
0)) # pad down
if norm == 'weight_norm':
self.conv = wn(self.conv)
elif norm == 'batch_norm':
self.bn = nn.BatchNorm2d(num_filters_out)
if shift_output_down:
self.down_shift = lambda x: down_shift(x, pad=nn.ZeroPad2d((0, 0, 1, 0)))
def __init__(self,
channel_in,
num_classes,
p_dropout=0.2,
weight_init=True):
super(DConvNet1, self).__init__()
self.logger = logging.getLogger(__name__) # initialize logger
self.num_classes = num_classes
# general reusable layers:
self.LReLU = nn.LeakyReLU(
negative_slope=0.2) # leaky ReLU activation function
self.sgmd = nn.Sigmoid() # sigmoid activation function
self.drop = nn.Dropout(p=p_dropout) # dropout layer
# input -->
# drop
# ConvConcat
self.conv1 = wn(
nn.Conv2d(in_channels=channel_in + num_classes,
out_channels=32,
kernel_size=(3, 3),
stride=(1, 1),
padding=1,
bias=False))
# LReLU
# ConvConcat
self.conv2 = wn(
nn.Conv2d(in_channels=32 + num_classes,
out_channels=32,
kernel_size=(3, 3),
stride=2,
padding=1,
bias=False))
# LReLU
# drop
# ConvConcat
self.conv3 = wn(
nn.Conv2d(in_channels=32 + num_classes,
out_channels=64,
kernel_size=(3, 3),
stride=(1, 1),
padding=1,
bias=False))
# LReLU
# ConvConcat
self.conv4 = wn(
nn.Conv2d(in_channels=64 + num_classes,
out_channels=64,
kernel_size=(3, 3),
stride=2,
padding=1,
bias=False))
# LReLU
# drop
# ConvConcat
self.conv5 = wn(
nn.Conv2d(in_channels=64 + num_classes,
out_channels=128,
kernel_size=(3, 3),
stride=(1, 1),
padding=0,
bias=False))
# LReLU
# ConvConcat
self.conv6 = wn(
nn.Conv2d(in_channels=128 + num_classes,
out_channels=128,
kernel_size=(3, 3),
stride=(1, 1),
padding=0,
bias=False))
# LReLU
self.globalPool = nn.AdaptiveAvgPool2d(output_size=4)
# MLPConcat
self.lin = nn.Linear(in_features=128 * 4 * 4 + num_classes,
out_features=1)
# smg
if weight_init:
# initialize weights for all conv and lin layers
self.apply(init_weights)
# log network structure
self.logger.debug(self)
def __init__(self, dim_in, dim_out):
super(nin, self).__init__()
self.lin_a = wn(
nn.Linear(dim_in, dim_out)
) # Nawid -Weight normalisation of a linear unit and the performs the linear transformation where the output dimension is dim_out
self.dim_out = dim_out # Nawid - Output dimension
def __init__(self, dim_in, dim_out):
super(nin, self).__init__()
self.lin_a = wn(nn.Linear(dim_in, dim_out))
self.dim_out = dim_out
def __init__(self,
nr_resnet=5,
nr_filters=80,
nr_logistic_mix=10,
resnet_nonlinearity='concat_elu',
input_channels=3,
kernel_size=(5, 5),
max_dilation=2,
weight_norm=True,
feature_norm_op=None,
dropout_prob=0.5,
conv_bias=True,
conv_mask_weight=False,
rematerialize=False,
binarize=False):
super(OurPixelCNN, self).__init__()
assert resnet_nonlinearity == 'concat_elu'
self.resnet_nonlinearity = lambda x: concat_elu(x)
self.init_padding = None
self.binarize = binarize
if weight_norm:
conv_op_init = lambda cin, cout: wn(
locally_masked_conv2d(cin,
cout,
kernel_size=kernel_size,
bias=conv_bias,
mask_weight=conv_mask_weight))
conv_op_dilated = lambda cin, cout: wn(
locally_masked_conv2d(cin,
cout,
kernel_size=kernel_size,
dilation=max_dilation,
bias=conv_bias,
mask_weight=conv_mask_weight))
conv_op = lambda cin, cout: wn(
locally_masked_conv2d(cin,
cout,
kernel_size=kernel_size,
bias=conv_bias,
mask_weight=conv_mask_weight))
else:
conv_op_init = lambda cin, cout: locally_masked_conv2d(
cin,
cout,
kernel_size=kernel_size,
bias=conv_bias,
mask_weight=conv_mask_weight)
conv_op_dilated = lambda cin, cout: locally_masked_conv2d(
cin,
cout,
kernel_size=kernel_size,
dilation=max_dilation,
bias=conv_bias,
mask_weight=conv_mask_weight)
conv_op = lambda cin, cout: locally_masked_conv2d(
cin,
cout,
kernel_size=kernel_size,
bias=conv_bias,
mask_weight=conv_mask_weight)
down_nr_resnet = [nr_resnet] + [nr_resnet + 1] * 2
self.down_layers = nn.ModuleList([
OurPixelCNNLayer_down(down_nr_resnet[i],
nr_filters,
self.resnet_nonlinearity,
conv_op,
feature_norm_op,
kernel_size=kernel_size,
weight_norm=weight_norm,
dropout_prob=dropout_prob,
rematerialize=rematerialize)
for i in range(3)
])
self.up_layers = nn.ModuleList([
OurPixelCNNLayer_up(nr_resnet,
nr_filters,
self.resnet_nonlinearity,
conv_op,
feature_norm_op,
kernel_size=kernel_size,
weight_norm=weight_norm,
dropout_prob=dropout_prob,
rematerialize=rematerialize) for _ in range(3)
])
self.u_init = conv_op_init(input_channels + 1, nr_filters)
self.downsize_u_stream = nn.ModuleList(
[conv_op_dilated(nr_filters, nr_filters) for _ in range(2)])
self.upsize_u_stream = nn.ModuleList(
[conv_op_dilated(nr_filters, nr_filters) for _ in range(2)])
self.norm_init = feature_norm_op(
nr_filters) if feature_norm_op else identity
self.norm_ds = nn.ModuleList(
[feature_norm_op(nr_filters)
for _ in range(2)]) if feature_norm_op else None
self.norm_us = nn.ModuleList(
[feature_norm_op(nr_filters)
for _ in range(2)]) if feature_norm_op else None
if self.binarize:
self.nin_out = nin(nr_filters, 2, weight_norm=True)
else:
num_mix = 3 if input_channels == 1 else 10
self.nin_out = nin(nr_filters,
num_mix * nr_logistic_mix,
weight_norm=True)
def __init__(self, params):
super(WNNet, self).__init__(params)
self.fc1 = wn(self.fc1) # add weight normalization here
self.__net__ = "WNNet"
def __init__(self, in_size, out_size):
super().__init__()
self.lin = wn(nn.Linear(in_size, out_size))
def __init__(self, dim_in, dim_out):
super(nin, self).__init__()
self.lin_a = wn(nn.Linear(dim_in, dim_out))
self.dim_out = dim_out
