def __init__(self, n_classes=6):
super(SteerCNN, self).__init__()
# the model is equivariant under rotations by 45 degrees, modelled by C8
self.r2_act = gspaces.Rot2dOnR2(N=4)
# the input image is a scalar field, corresponding to the trivial representation
input_type = nn_e2.FieldType(self.r2_act,
3 * [self.r2_act.trivial_repr])
# we store the input type for wrapping the images into a geometric tensor during the forward pass
self.input_type = input_type
# convolution 1
# first specify the output type of the convolutional layer
# we choose 24 feature fields, each transforming under the regular representation of C8
out_type = nn_e2.FieldType(self.r2_act,
24 * [self.r2_act.regular_repr])
self.block1 = nn_e2.SequentialModule(
nn_e2.R2Conv(input_type,
out_type,
kernel_size=7,
padding=3,
bias=False), nn_e2.InnerBatchNorm(out_type),
nn_e2.ReLU(out_type, inplace=True))
self.pool1 = nn_e2.PointwiseAvgPool(out_type, 4)
# convolution 2
# the old output type is the input type to the next layer
in_type = self.block1.out_type
# the output type of the second convolution layer are 48 regular feature fields of C8
#out_type = nn_e2.FieldType(self.r2_act, 48 * [self.r2_act.regular_repr])
self.block2 = nn_e2.SequentialModule(
nn_e2.R2Conv(in_type,
out_type,
kernel_size=7,
padding=3,
bias=False), nn_e2.InnerBatchNorm(out_type),
nn_e2.ReLU(out_type, inplace=True))
self.pool2 = nn_e2.SequentialModule(
nn_e2.PointwiseAvgPoolAntialiased(out_type,
sigma=0.66,
stride=1,
padding=0),
nn_e2.PointwiseAvgPool(out_type, 4), nn_e2.GroupPooling(out_type))
# PointwiseAvgPoolAntialiased(out_type, sigma=0.66, stride=7)
# number of output channels
c = 24 * 13 * 13 #self.gpool.out_type.size
# Fully Connected
self.fully_net = torch.nn.Sequential(
torch.nn.Linear(c, 64),
torch.nn.BatchNorm1d(64),
torch.nn.ELU(inplace=True),
torch.nn.Linear(64, n_classes),
)
def build_multiscale_classifier(self, input_size):
n, c, h, w = input_size
hidden_shapes = []
for i in range(self.n_scale):
if i < self.n_scale - 1:
c *= 2 if self.factor_out else 4
h //= 2
w //= 2
hidden_shapes.append((n, c, h, w))
classification_heads = []
feat_type_out = FIBERS['regular'](self.group_action_type,
self.classification_hdim,
self.field_type, fixparams=True)
feat_type_mid = FIBERS['regular'](self.group_action_type,
int(self.classification_hdim // 2),
self.field_type, fixparams=True)
feat_type_last = FIBERS['regular'](self.group_action_type,
int(self.classification_hdim // 4),
self.field_type, fixparams=True)
# feat_type_out = enn.FieldType(self.group_action_type,
# self.classification_hdim*[self.group_action_type.regular_repr])
for i, hshape in enumerate(hidden_shapes):
classification_heads.append(
nn.Sequential(
enn.R2Conv(self.input_type, feat_type_out, 5, stride=2),
layers.EquivariantActNorm2d(feat_type_out.size),
enn.ReLU(feat_type_out, inplace=True),
enn.PointwiseAvgPoolAntialiased(feat_type_out, sigma=0.66, stride=2),
enn.R2Conv(feat_type_out, feat_type_mid, kernel_size=3),
layers.EquivariantActNorm2d(feat_type_mid.size),
enn.ReLU(feat_type_mid, inplace=True),
enn.PointwiseAvgPoolAntialiased(feat_type_mid, sigma=0.66, stride=1),
enn.R2Conv(feat_type_mid, feat_type_last, kernel_size=3),
layers.EquivariantActNorm2d(feat_type_last.size),
enn.ReLU(feat_type_last, inplace=True),
enn.PointwiseAvgPoolAntialiased(feat_type_last, sigma=0.66, stride=2),
enn.GroupPooling(feat_type_last),
)
)
self.classification_heads = nn.ModuleList(classification_heads)
self.logit_layer = nn.Linear(classification_heads[-1][-1].out_type.size, self.n_classes)
def __init__(self, input_size, in_type, field_type, out_fiber,
activation_fn, hidden_size, group_action_type):
super(InvariantCNNBlock, self).__init__()
_, self.c, self.h, self.w = input_size
ngf = 16
self.group_action_type = group_action_type
feat_type_in = enn.FieldType(
self.group_action_type,
self.c * [self.group_action_type.trivial_repr])
feat_type_hid = FIBERS[out_fiber](group_action_type,
hidden_size,
field_type,
fixparams=True)
feat_type_out = enn.FieldType(
self.group_action_type,
128 * [self.group_action_type.regular_repr])
# we store the input type for wrapping the images into a geometric tensor during the forward pass
self.input_type = feat_type_in
self.block1 = enn.SequentialModule(
enn.R2Conv(feat_type_in, feat_type_hid, kernel_size=5, padding=0),
enn.InnerBatchNorm(feat_type_hid),
activation_fn(feat_type_hid, inplace=True),
)
self.pool1 = enn.SequentialModule(
enn.PointwiseAvgPoolAntialiased(feat_type_hid,
sigma=0.66,
stride=2))
self.block2 = enn.SequentialModule(
enn.R2Conv(feat_type_hid, feat_type_hid, kernel_size=5),
enn.InnerBatchNorm(feat_type_hid),
activation_fn(feat_type_hid, inplace=True),
)
self.pool2 = enn.SequentialModule(
enn.PointwiseAvgPoolAntialiased(feat_type_hid,
sigma=0.66,
stride=2))
self.block3 = enn.SequentialModule(
enn.R2Conv(feat_type_hid, feat_type_out, kernel_size=3, padding=1),
enn.InnerBatchNorm(feat_type_out),
activation_fn(feat_type_out, inplace=True),
)
self.pool3 = enn.PointwiseAvgPoolAntialiased(feat_type_out,
sigma=0.66,
stride=1,
padding=0)
self.gpool = enn.GroupPooling(feat_type_out)
self.gc = self.gpool.out_type.size
self.gen = torch.nn.Sequential(
torch.nn.ConvTranspose2d(self.gc,
ngf,
kernel_size=4,
stride=1,
padding=0), torch.nn.BatchNorm2d(ngf),
torch.nn.LeakyReLU(0.2),
torch.nn.ConvTranspose2d(ngf,
ngf,
kernel_size=4,
stride=2,
padding=1), torch.nn.BatchNorm2d(ngf),
torch.nn.LeakyReLU(0.2),
torch.nn.ConvTranspose2d(ngf,
int(ngf / 2),
kernel_size=4,
stride=2,
padding=1),
torch.nn.BatchNorm2d(int(ngf / 2)), torch.nn.LeakyReLU(0.2),
torch.nn.ConvTranspose2d(int(ngf / 2),
self.c,
kernel_size=4,
stride=2,
padding=1), torch.nn.Tanh())
def __init__(self, n_classes=10):
super(C8SteerableCNN, self).__init__()
# the model is equivariant under rotations by 45 degrees, modelled by C8
self.r2_act = gspaces.Rot2dOnR2(N=8)
# the input image is a scalar field, corresponding to the trivial representation
in_type = nn.FieldType(self.r2_act, [self.r2_act.trivial_repr])
# we store the input type for wrapping the images into a geometric tensor during the forward pass
self.input_type = in_type
# convolution 1
# first specify the output type of the convolutional layer
# we choose 16 feature fields, each transforming under the regular representation of C8
out_type = nn.FieldType(self.r2_act, 24 * [self.r2_act.regular_repr])
self.block1 = nn.SequentialModule(
# nn.MaskModule(in_type, 29, margin=1),
nn.R2Conv(in_type, out_type, kernel_size=7, padding=1, bias=False),
nn.InnerBatchNorm(out_type),
nn.ReLU(out_type, inplace=True))
# convolution 2
# the old output type is the input type to the next layer
in_type = self.block1.out_type
# the output type of the second convolution layer are 32 regular feature fields of C8
out_type = nn.FieldType(self.r2_act, 48 * [self.r2_act.regular_repr])
self.block2 = nn.SequentialModule(
nn.R2Conv(in_type, out_type, kernel_size=5, padding=2, bias=False),
nn.InnerBatchNorm(out_type), nn.ReLU(out_type, inplace=True))
self.pool1 = nn.SequentialModule(
nn.PointwiseAvgPoolAntialiased(out_type, sigma=0.66, stride=2))
# convolution 3
# the old output type is the input type to the next layer
in_type = self.block2.out_type
# the output type of the third convolution layer are 32 regular feature fields of C8
out_type = nn.FieldType(self.r2_act, 48 * [self.r2_act.regular_repr])
self.block3 = nn.SequentialModule(
nn.R2Conv(in_type, out_type, kernel_size=5, padding=2, bias=False),
nn.InnerBatchNorm(out_type), nn.ReLU(out_type, inplace=True))
# convolution 4
# the old output type is the input type to the next layer
in_type = self.block3.out_type
# the output type of the fourth convolution layer are 64 regular feature fields of C8
out_type = nn.FieldType(self.r2_act, 96 * [self.r2_act.regular_repr])
self.block4 = nn.SequentialModule(
nn.R2Conv(in_type, out_type, kernel_size=5, padding=2, bias=False),
nn.InnerBatchNorm(out_type), nn.ReLU(out_type, inplace=True))
self.pool2 = nn.SequentialModule(
nn.PointwiseAvgPoolAntialiased(out_type, sigma=0.66, stride=2))
# convolution 5
# the old output type is the input type to the next layer
in_type = self.block4.out_type
# the output type of the fifth convolution layer are 64 regular feature fields of C8
out_type = nn.FieldType(self.r2_act, 96 * [self.r2_act.regular_repr])
self.block5 = nn.SequentialModule(
nn.R2Conv(in_type, out_type, kernel_size=5, padding=2, bias=False),
nn.InnerBatchNorm(out_type), nn.ReLU(out_type, inplace=True))
# convolution 6
# the old output type is the input type to the next layer
in_type = self.block5.out_type
# the output type of the sixth convolution layer are 64 regular feature fields of C8
out_type = nn.FieldType(self.r2_act, 64 * [self.r2_act.regular_repr])
self.block6 = nn.SequentialModule(
nn.R2Conv(in_type, out_type, kernel_size=5, padding=1, bias=False),
nn.InnerBatchNorm(out_type), nn.ReLU(out_type, inplace=True))
self.pool3 = nn.PointwiseAvgPool(out_type, kernel_size=4)
self.gpool = nn.GroupPooling(out_type)
# number of output channels
c = self.gpool.out_type.size
# Fully Connected
self.fully_net = torch.nn.Sequential(
torch.nn.Linear(c, 64),
torch.nn.BatchNorm1d(64),
torch.nn.ELU(inplace=True),
torch.nn.Linear(64, n_classes),
)
def __init__(self, input_shape, num_actions, dueling_DQN, rest_lev):
super(steerable_DQN_Pacman, self).__init__()
self.input_shape = input_shape
self.num_actions = num_actions
self.dueling_DQN = dueling_DQN
self.rest_lev = rest_lev
self.feature_factor = 1
# Scales up the num of fields
self.num_feature_fields = [8, 16, 16, 128]
# Symmetry group
self.r2_act = gspaces.FlipRot2dOnR2(N=4)
# 1st E-Conv
self.input_type = nn.FieldType(
self.r2_act, input_shape[0] * [self.r2_act.trivial_repr])
feature1_type = nn.FieldType(
self.r2_act,
self.num_feature_fields[0] * [self.r2_act.regular_repr])
self.feature_field1 = nn.SequentialModule(
nn.R2Conv(self.input_type,
feature1_type,
kernel_size=7,
padding=2,
stride=2,
bias=False), nn.ReLU(feature1_type, inplace=True),
nn.PointwiseAvgPoolAntialiased(feature1_type, sigma=0.66,
stride=2))
# 2nd E-Conv
self.input_feature_type = feature1_type
if self.rest_lev == 3:
self.restrict1 = self.restrict_layer((0, 2), feature1_type)
self.feature_factor = 2
else:
self.restrict1 = lambda x: x
feature2_type = nn.FieldType(
self.r2_act, self.num_feature_fields[1] * self.feature_factor *
[self.r2_act.regular_repr])
self.feature_field2 = nn.SequentialModule(
nn.R2Conv(self.input_feature_type,
feature2_type,
kernel_size=5,
padding=2,
stride=2,
bias=False), nn.ReLU(feature2_type, inplace=True))
# 3rd E-Conv
self.input_feature_type = feature2_type
if self.rest_lev == 2:
self.restrict2 = self.restrict_layer((0, 1), feature2_type)
self.feature_factor = 4
else:
self.restrict2 = lambda x: x
feature3_type = nn.FieldType(
self.r2_act, self.num_feature_fields[2] * self.feature_factor *
[self.r2_act.regular_repr])
self.feature_field3 = nn.SequentialModule(
nn.R2Conv(self.input_feature_type,
feature3_type,
kernel_size=5,
padding=1,
stride=2,
bias=False), nn.ReLU(feature3_type, inplace=True))
# 4th E-Conv
self.input_feature_type = feature3_type
if rest_lev == 1:
self.restrict_extra = self.restrict_layer((0, 1), feature3_type)
self.feature_factor = 4
else:
self.restrict_extra = lambda x: x
feature4_type = nn.FieldType(
self.r2_act, self.num_feature_fields[3] * self.feature_factor *
[self.r2_act.regular_repr])
self.feature_field_extra = nn.SequentialModule(
nn.R2Conv(self.input_feature_type,
feature4_type,
kernel_size=5,
padding=0,
bias=True), nn.ReLU(feature4_type, inplace=True))
_, _ = self.feature_shape()
self.out_size = self.feature_field_extra.out_type.size
# Final linear layer
if self.dueling_DQN:
print("You are using Dueling DQN")
self.advantage = torch.nn.Linear(self.out_size, self.num_actions)
self.value = torch.nn.Linear(self.out_size, 1)
else:
self.actionvalue = torch.nn.Linear(self.out_size, self.num_actions)