def __init__(self, in_channels=12, num_action=4):
super(DQN_Dueling, self).__init__()
self.num_action = num_action
self.conv1 = nn.Conv2d(in_channels, 32, kernel_size=8, stride=4)
self.conv2 = nn.Conv2d(32, 64, kernel_size=4, stride=2)
self.conv3 = nn.Conv2d(64, 64, kernel_size=3, stride=1)
self.conv2 = AugmentedConv(in_channels=32,
out_channels=64,
kernel_size=4,
dk=16,
dv=4,
Nh=4,
stride=2)
self.conv3 = AugmentedConv(in_channels=64,
out_channels=64,
kernel_size=3,
dk=8,
dv=4,
Nh=4,
stride=1)
self.A1 = nn.Linear(10 * 10 * 64, 512)
self.V1 = nn.Linear(10 * 10 * 64, 512)
self.A2 = nn.Linear(512, num_action)
self.V2 = nn.Linear(512, 1)
def __init__(self, in_planes, planes, dropout_rate, stride=1, v=0.2, k=2, Nh=4):
super(wide_basic, self).__init__()
self.bn1 = nn.BatchNorm2d(in_planes)
self.conv1 = AugmentedConv(in_planes, planes, kernel_size=3, dk=k * planes, dv=int(v * planes), Nh=Nh, relative=True, padding=1)
self.dropout = nn.Dropout(p=dropout_rate)
self.bn2 = nn.BatchNorm2d(planes)
self.conv2 = AugmentedConv(planes, planes, kernel_size=3, dk=k * planes, dv=int(v * planes), Nh=Nh, stride=stride, relative=True, padding=1)
self.shortcut = nn.Sequential()
if stride != 1 or in_planes != planes:
self.shortcut = nn.Sequential(
AugmentedConv(in_planes, planes, kernel_size=3, dk=k * planes, dv=int(v * planes), Nh=Nh, relative=True, stride=stride, padding=1),
)
def __init__(self, depth, widen_factor, dropout_rate, num_classes):
super(Wide_ResNet, self).__init__()
self.in_planes = 20
assert ((depth-4) % 6 == 0), 'Wide-resnet depth should be 6n+4'
n = int((depth - 4) / 6)
k = widen_factor
dv_v = 0.2
dk_k = 2
Nh = 4
print('| Wide-Resnet %dx%d' % (depth, k))
n_Stages = [20, 20 * k, 40 * k, 40 * k]
self.conv1 = AugmentedConv(3, n_Stages[0], kernel_size=3, dk=dk_k * n_Stages[0], dv=int(dv_v * n_Stages[0]), Nh=Nh, relative=True, padding=1)
self.layer1 = nn.Sequential(
self._wide_layer(wide_basic, n_Stages[1], n, dropout_rate, stride=1),
)
self.layer2 = nn.Sequential(
self._wide_layer(wide_basic, n_Stages[2], n, dropout_rate, stride=2),
)
self.layer3 = nn.Sequential(
self._wide_layer(wide_basic, n_Stages[3], n, dropout_rate, stride=2),
)
self.bn1 = nn.BatchNorm2d(n_Stages[3], momentum=0.9)
self.linear = nn.Linear(n_Stages[3], num_classes)
def __init__(self,
in_channels,
out_channels,
stride=1,
base_width=64,
args=None):
super(Bottleneck, self).__init__()
self.stride = stride
groups = args.groups # Number of attention heads
'''
# TODO : Doubt in width, when base_width != 64?
width = int(out_channels * (base_width / 64.))\
if args.attention_conv\
else int(out_channels * (base_width / 64.)) * groups
'''
width = out_channels
additional_args = {'groups':groups, 'R':args.R, 'z_init':args.z_init, 'adaptive_span':args.adaptive_span} \
if args.all_attention else {'bias': False}
kernel_size = args.attention_kernel if args.all_attention else 3
padding = int((kernel_size - 1) / 2)
layer = None
if args.attention_conv:
layer = AugmentedConv(width,
width,
kernel_size,
args.dk,
args.dv,
groups,
shape=width)
elif args.all_attention:
layer = AttentionConv(width,
width,
kernel_size=kernel_size,
padding=padding,
**additional_args)
else:
layer = nn.Conv2d(width,
width,
kernel_size=kernel_size,
padding=padding,
**additional_args)
self.conv1 = nn.Sequential(
nn.Conv2d(in_channels, width, kernel_size=1, bias=False),
nn.BatchNorm2d(width),
nn.ReLU(),
)
self.conv2 = nn.Sequential(
layer,
nn.BatchNorm2d(width),
nn.ReLU(),
)
self.conv3 = nn.Sequential(
nn.Conv2d(width,
self.expansion * out_channels,
kernel_size=1,
bias=False),
nn.BatchNorm2d(self.expansion * out_channels),
)
self.shortcut = nn.Sequential()
if stride != 1 or in_channels != self.expansion * out_channels:
self.shortcut = nn.Sequential(
nn.Conv2d(in_channels,
self.expansion * out_channels,
kernel_size=1,
stride=stride,
bias=False),
nn.BatchNorm2d(self.expansion * out_channels))