def __init__(self, input_channels, hidden_dim, kernel_size, N # Group size ): super(rot_resblock, self).__init__() # Specify symmetry transformation r2_act = gspaces.Rot2dOnR2(N = N) feat_type_in = nn.FieldType(r2_act, input_channels*[r2_act.regular_repr]) feat_type_hid = nn.FieldType(r2_act, hidden_dim*[r2_act.regular_repr]) self.layer1 = nn.SequentialModule( nn.R2Conv(feat_type_in, feat_type_hid, kernel_size = kernel_size, padding = (kernel_size - 1)//2), nn.InnerBatchNorm(feat_type_hid), nn.ReLU(feat_type_hid) ) self.layer2 = nn.SequentialModule( nn.R2Conv(feat_type_hid, feat_type_hid, kernel_size = kernel_size, padding = (kernel_size - 1)//2), nn.InnerBatchNorm(feat_type_hid), nn.ReLU(feat_type_hid) ) self.upscale = nn.SequentialModule( nn.R2Conv(feat_type_in, feat_type_hid, kernel_size = kernel_size, padding = (kernel_size - 1)//2), nn.InnerBatchNorm(feat_type_hid), nn.ReLU(feat_type_hid) ) self.input_channels = input_channels self.hidden_dim = hidden_dim
def __init__(self, input_channels, output_channels, kernel_size, stride, N, activation = True, deconv = False, last_deconv = False): super(rot_conv2d, self).__init__() r2_act = gspaces.Rot2dOnR2(N = N) feat_type_in = nn.FieldType(r2_act, input_channels*[r2_act.regular_repr]) feat_type_hid = nn.FieldType(r2_act, output_channels*[r2_act.regular_repr]) if not deconv: if activation: self.layer = nn.SequentialModule( nn.R2Conv(feat_type_in, feat_type_hid, kernel_size = kernel_size, stride = stride, padding = (kernel_size - 1)//2), nn.InnerBatchNorm(feat_type_hid), nn.ReLU(feat_type_hid) ) else: self.layer = nn.R2Conv(feat_type_in, feat_type_hid, kernel_size = kernel_size, stride = stride,padding = (kernel_size - 1)//2) else: if last_deconv: feat_type_in = nn.FieldType(r2_act, input_channels*[r2_act.regular_repr]) feat_type_hid = nn.FieldType(r2_act, output_channels*[r2_act.irrep(1)]) self.layer = nn.R2Conv(feat_type_in, feat_type_hid, kernel_size = kernel_size, stride = stride, padding = 0) else: self.layer = nn.SequentialModule( nn.R2Conv(feat_type_in, feat_type_hid, kernel_size = kernel_size, stride = stride, padding = 0), nn.InnerBatchNorm(feat_type_hid), nn.ReLU(feat_type_hid) )
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 __init__(self): super(DenseFeatureExtractionModuleE2Inv, self).__init__() filters = np.array([32,32, 64,64, 128,128,128, 256,256,256, 512,512,512], dtype=np.int32)*2 # number of rotations to consider for rotation invariance N = 8 self.gspace = gspaces.Rot2dOnR2(N) self.input_type = enn.FieldType(self.gspace, [self.gspace.trivial_repr] * 3) ip_op_types = [ self.input_type, ] self.num_channels = 64 for filter_ in filters[:10]: ip_op_types.append(FIELD_TYPE['regular'](self.gspace, filter_, fixparams=False)) self.model = enn.SequentialModule(*[ conv3x3(ip_op_types[0], ip_op_types[1]), enn.ReLU(ip_op_types[1], inplace=True), conv3x3(ip_op_types[1], ip_op_types[2]), enn.ReLU(ip_op_types[2], inplace=True), enn.PointwiseMaxPool(ip_op_types[2], 2), conv3x3(ip_op_types[2], ip_op_types[3]), enn.ReLU(ip_op_types[3], inplace=True), conv3x3(ip_op_types[3], ip_op_types[4]), enn.ReLU(ip_op_types[4], inplace=True), enn.PointwiseMaxPool(ip_op_types[4], 2), conv3x3(ip_op_types[4], ip_op_types[5]), enn.ReLU(ip_op_types[5], inplace=True), conv3x3(ip_op_types[5], ip_op_types[6]), enn.ReLU(ip_op_types[6], inplace=True), conv3x3(ip_op_types[6], ip_op_types[7]), enn.ReLU(ip_op_types[7], inplace=True), enn.PointwiseAvgPool(ip_op_types[7], kernel_size=2, stride=1), conv5x5(ip_op_types[7], ip_op_types[8]), enn.ReLU(ip_op_types[8], inplace=True), conv5x5(ip_op_types[8], ip_op_types[9]), enn.ReLU(ip_op_types[9], inplace=True), conv5x5(ip_op_types[9], ip_op_types[10]), enn.ReLU(ip_op_types[10], inplace=True), # enn.PointwiseMaxPool(ip_op_types[7], 2), # conv3x3(ip_op_types[7], ip_op_types[8]), # enn.ReLU(ip_op_types[8], inplace=True), # conv3x3(ip_op_types[8], ip_op_types[9]), # enn.ReLU(ip_op_types[9], inplace=True), # conv3x3(ip_op_types[9], ip_op_types[10]), # enn.ReLU(ip_op_types[10], inplace=True), enn.GroupPooling(ip_op_types[10]) ])
def __init__(self, in_type, inner_type, out_type, stride=1): super(BottleneckBlock, self).__init__() self.bn1 = enn.InnerBatchNorm(in_type) self.relu1 = enn.ReLU(in_type,inplace=True) self.conv1 = conv1x1(in_type, inner_type) self.bn2 = enn.InnerBatchNorm(inner_type) self.relu2 = enn.ReLU(inner_type,inplace=True) self.conv2 = conv3x3(inner_type, out_type)
def __init__(self, input_shape, num_actions, dueling_DQN): super(D4_steerable_DQN_Snake, self).__init__() self.input_shape = input_shape self.num_actions = num_actions self.dueling_DQN = dueling_DQN self.r2_act = gspaces.FlipRot2dOnR2(N=4) self.input_type = nn.FieldType( self.r2_act, input_shape[0] * [self.r2_act.trivial_repr]) feature1_type = nn.FieldType(self.r2_act, 8 * [self.r2_act.regular_repr]) feature2_type = nn.FieldType(self.r2_act, 12 * [self.r2_act.regular_repr]) feature3_type = nn.FieldType(self.r2_act, 12 * [self.r2_act.regular_repr]) feature4_type = nn.FieldType(self.r2_act, 32 * [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)) self.feature_field2 = nn.SequentialModule( nn.R2Conv(feature1_type, feature2_type, kernel_size=5, padding=1, stride=2, bias=False), nn.ReLU(feature2_type, inplace=True)) self.feature_field3 = nn.SequentialModule( nn.R2Conv(feature2_type, feature3_type, kernel_size=5, padding=1, stride=1, bias=False), nn.ReLU(feature3_type, inplace=True)) self.equivariant_features = nn.SequentialModule( nn.R2Conv(feature3_type, feature4_type, kernel_size=5, stride=1, bias=False), nn.ReLU(feature4_type, inplace=True)) self.gpool = nn.GroupPooling(feature4_type) self.feature_shape() if self.dueling_DQN: print("You are using Dueling DQN") self.advantage = torch.nn.Linear( self.equivariant_features.out_type.size, self.num_actions) #self.value = torch.nn.Linear(self.gpool.out_type.size, 1) self.value = torch.nn.Linear( self.equivariant_features.out_type.size, 1) else: self.actionvalue = torch.nn.Linear( self.equivariant_features.out_type.size, self.num_actions)
def __init__(self, base='DNSteerableLeNet', in_chan=1, n_classes=2, imsize=150, kernel_size=5, N=8, quiet=True, number_rotations=None): super(DNSteerableLeNet, self).__init__() kernel_size = int(kernel_size) out_chan = int(n_classes) if number_rotations != None: N = int(number_rotations) z = imsize // 2 // 2 self.r2_act = gspaces.FlipRot2dOnR2(N) in_type = e2nn.FieldType(self.r2_act, [self.r2_act.trivial_repr]) self.input_type = in_type out_type = e2nn.FieldType(self.r2_act, 6 * [self.r2_act.regular_repr]) self.mask = e2nn.MaskModule(in_type, imsize, margin=1) self.conv1 = e2nn.R2Conv(in_type, out_type, kernel_size=kernel_size, padding=kernel_size // 2, bias=False) self.relu1 = e2nn.ReLU(out_type, inplace=True) self.pool1 = e2nn.PointwiseMaxPoolAntialiased(out_type, kernel_size=2) in_type = self.pool1.out_type out_type = e2nn.FieldType(self.r2_act, 16 * [self.r2_act.regular_repr]) self.conv2 = e2nn.R2Conv(in_type, out_type, kernel_size=kernel_size, padding=kernel_size // 2, bias=False) self.relu2 = e2nn.ReLU(out_type, inplace=True) self.pool2 = e2nn.PointwiseMaxPoolAntialiased(out_type, kernel_size=2) self.gpool = e2nn.GroupPooling(out_type) self.fc1 = nn.Linear(16 * z * z, 120) self.fc2 = nn.Linear(120, 84) self.fc3 = nn.Linear(84, out_chan) self.drop = nn.Dropout(p=0.5) # dummy parameter for tracking device self.dummy = nn.Parameter(torch.empty(0))
def __init__(self, in_type, conv_func): super(RegressionHead, self).__init__() gspace = in_type.gspace self.add_module('gpool', nn.PointwiseAdaptiveMaxPool(in_type, (1, 1))) if isinstance(in_type.gspace, e2cnn.gspaces.Rot2dOnR2): base = 8 elif isinstance(in_type.gspace, e2cnn.gspaces.FlipRot2dOnR2): base = 4 # number of output channels # Fully Connected in_type = in_type out_type = nn.FieldType(gspace, 2 * base * [gspace.regular_repr]) self.add_module( 'block1', nn.SequentialModule( conv_func(in_type, out_type, kernel_size=1, padding=0, bias=False), nn.ReLU(out_type, inplace=True))) in_type = out_type if isinstance(gspace, gspaces.Rot2dOnR2): out_type = nn.FieldType(gspace, [gspace.irrep(1)]) elif isinstance(gspace, gspaces.FlipRot2dOnR2): out_type = nn.FieldType(gspace, [gspace.irrep(1, 1)]) else: raise NotImplementedError self.add_module( 'block2', conv_func(in_type, out_type, kernel_size=1, padding=0, bias=False))
def __init__(self, channel_in=3, n_classes=4, rot_n=4): super(SmallE2, self).__init__() r2_act = gspaces.Rot2dOnR2(N=rot_n) self.feat_type_in = nn.FieldType(r2_act, channel_in * [r2_act.trivial_repr]) feat_type_hid = nn.FieldType(r2_act, 8 * [r2_act.regular_repr]) feat_type_out = nn.FieldType(r2_act, 2 * [r2_act.regular_repr]) self.bn = nn.InnerBatchNorm(feat_type_hid) self.relu = nn.ReLU(feat_type_hid) self.convin = nn.R2Conv(self.feat_type_in, feat_type_hid, kernel_size=3) self.convhid = nn.R2Conv(feat_type_hid, feat_type_hid, kernel_size=3) self.convout = nn.R2Conv(feat_type_hid, feat_type_out, kernel_size=3) self.avgpool = nn.PointwiseAvgPool(feat_type_out, 3) self.invariant_map = nn.GroupPooling(feat_type_out) c = self.invariant_map.out_type.size self.lin_in = torch.nn.Linear(c, 64) self.elu = torch.nn.ELU() self.lin_out = torch.nn.Linear(64, n_classes)
def __init__(self, input_frames, output_frames, kernel_size, N): super(Unet_Rot, self).__init__() r2_act = gspaces.Rot2dOnR2(N = N) self.feat_type_in = nn.FieldType(r2_act, input_frames*[r2_act.irrep(1)]) self.feat_type_in_hid = nn.FieldType(r2_act, 32*[r2_act.regular_repr]) self.feat_type_hid_out = nn.FieldType(r2_act, (16 + input_frames)*[r2_act.irrep(1)]) self.feat_type_out = nn.FieldType(r2_act, output_frames*[r2_act.irrep(1)]) self.conv1 = nn.SequentialModule( nn.R2Conv(self.feat_type_in, self.feat_type_in_hid, kernel_size = kernel_size, stride = 2, padding = (kernel_size - 1)//2), nn.InnerBatchNorm(self.feat_type_in_hid), nn.ReLU(self.feat_type_in_hid) ) self.conv2 = rot_conv2d(32, 64, kernel_size = kernel_size, stride = 1, N = N) self.conv2_1 = rot_conv2d(64, 64, kernel_size = kernel_size, stride = 1, N = N) self.conv3 = rot_conv2d(64, 128, kernel_size = kernel_size, stride = 2, N = N) self.conv3_1 = rot_conv2d(128, 128, kernel_size = kernel_size, stride = 1, N = N) self.conv4 = rot_conv2d(128, 256, kernel_size = kernel_size, stride = 2, N = N) self.conv4_1 = rot_conv2d(256, 256, kernel_size = kernel_size, stride = 1, N = N) self.deconv3 = rot_deconv2d(256, 64, N) self.deconv2 = rot_deconv2d(192, 32, N) self.deconv1 = rot_deconv2d(96, 16, N, last_deconv = True) self.output_layer = nn.R2Conv(self.feat_type_hid_out, self.feat_type_out, kernel_size = kernel_size, padding = (kernel_size - 1)//2)
def __init__(self, in_chan, out_chan, imsize, kernel_size=5, N=8): super(DNRestrictedLeNet, self).__init__() z = imsize // 2 // 2 self.r2_act = gspaces.FlipRot2dOnR2(N) in_type = e2nn.FieldType(self.r2_act, [self.r2_act.trivial_repr]) self.input_type = in_type out_type = e2nn.FieldType(self.r2_act, 6 * [self.r2_act.regular_repr]) self.mask = e2nn.MaskModule(in_type, imsize, margin=1) self.conv1 = e2nn.R2Conv(in_type, out_type, kernel_size=kernel_size, padding=kernel_size // 2, bias=False) self.relu1 = e2nn.ReLU(out_type, inplace=True) self.pool1 = e2nn.PointwiseMaxPoolAntialiased(out_type, kernel_size=2) self.gpool = e2nn.GroupPooling(out_type) self.conv2 = nn.Conv2d(6, 16, kernel_size, padding=kernel_size // 2) self.fc1 = nn.Linear(16 * z * z, 120) self.fc2 = nn.Linear(120, 84) self.fc3 = nn.Linear(84, out_chan) self.drop = nn.Dropout(p=0.5) # dummy parameter for tracking device self.dummy = nn.Parameter(torch.empty(0))
def __init__(self, in_type, num_classes=10): super(ClassificationHead, self).__init__() gspace = in_type.gspace self.add_module('gpool', nn.GroupPooling(in_type)) # number of output channels # Fully Connected in_type = self.gpool.out_type out_type = nn.FieldType(gspace, 64 * [gspace.trivial_repr]) self.add_module( 'linear1', sscnn.e2cnn.PlainConv(in_type, out_type, kernel_size=1, padding=0, bias=False)) self.add_module('relu1', nn.ReLU(out_type, inplace=True)) in_type = out_type out_type = nn.FieldType(gspace, num_classes * [gspace.trivial_repr]) self.add_module( 'linear2', sscnn.e2cnn.PlainConv(in_type, out_type, kernel_size=1, padding=0, bias=False))
def __init__( self, in_type: enn.FieldType, inner_type: enn.FieldType, dropout_rate: float, stride: int = 1, out_type: enn.FieldType = None, ): super(WideBasic, self).__init__() if out_type is None: out_type = in_type self.in_type = in_type inner_type = inner_type self.out_type = out_type if isinstance(in_type.gspace, gspaces.FlipRot2dOnR2): rotations = in_type.gspace.fibergroup.rotation_order elif isinstance(in_type.gspace, gspaces.Rot2dOnR2): rotations = in_type.gspace.fibergroup.order() else: rotations = 0 if rotations in [0, 2, 4]: conv = conv3x3 else: conv = conv5x5 self.bn1 = enn.InnerBatchNorm(self.in_type) self.relu1 = enn.ReLU(self.in_type, inplace=True) self.conv1 = conv(self.in_type, inner_type) self.bn2 = enn.InnerBatchNorm(inner_type) self.relu2 = enn.ReLU(inner_type, inplace=True) self.dropout = enn.PointwiseDropout(inner_type, p=dropout_rate) self.conv2 = conv(inner_type, self.out_type, stride=stride) self.shortcut = None if stride != 1 or self.in_type != self.out_type: self.shortcut = conv1x1(self.in_type, self.out_type, stride=stride, bias=False)
def __init__(self, in_type, out_type): super(BasicBlock, self).__init__() self.in_type = in_type self.out_type = out_type self.bn1 = enn.InnerBatchNorm(self.in_type) self.relu1 = enn.ReLU(self.in_type, inplace=True) self.conv1 = conv3x3(self.in_type, self.out_type)
def __init__(self, in_type, out_type, gspace): super(TransitionBlock, self).__init__() self.gspace = gspace self.in_type = FIELD_TYPE["regular"](self.gspace, in_type, fixparams=False) self.out_type = FIELD_TYPE["regular"](self.gspace, out_type, fixparams=False) self.bn1 = enn.InnerBatchNorm(self.in_type) self.relu1 = enn.ReLU(self.in_type,inplace=True) self.conv1 = conv1x1(self.in_type,self.out_type) self.avgpool = enn.PointwiseAvgPool(self.out_type, kernel_size=2)
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, in_type,inner_type, out_type, stride=1): super(ResBlock, self).__init__() self.in_type = in_type self.inner_type = inner_type self.out_type = out_type self.conv1 = conv1x1(self.in_type, self.inner_type, stride = 1, bias = False) self.bn1 = enn.InnerBatchNorm(self.inner_type) self.relu1 = enn.ReLU(self.inner_type) self.conv2 = conv3x3(self.inner_type, self.inner_type, padding=1, stride = stride, bias = False) self.bn2 = enn.InnerBatchNorm(self.inner_type) self.relu2 = enn.ReLU(self.inner_type, inplace=True) self.conv3 = conv1x1(self.inner_type, self.out_type, stride = 1, bias = False) self.bn3 = enn.InnerBatchNorm(self.out_type) self.relu3 = enn.ReLU(self.out_type, inplace=True) self.shortcut = None if stride != 1 or self.in_type != self.out_type: self.shortcut = enn.R2Conv(self.in_type, self.out_type, kernel_size=1, stride=stride, bias=False)
def __init__(self, growth_rate, list_layer, nclasses): super(DenseNet161, self).__init__() self.gspace = gspaces.Rot2dOnR2(N=8) in_type = 2*growth_rate self.conv1 = conv7x7(FIELD_TYPE["trivial"](self.gspace, 3, fixparams=False), FIELD_TYPE["regular"](self.gspace, in_type, fixparams=False)) self.pool1 = enn.PointwiseMaxPool(FIELD_TYPE["regular"](self.gspace, in_type, fixparams=False), kernel_size=2, stride=2) #1st block self.block1 = DenseBlock(in_type, growth_rate, self.gspace, list_layer[0]) in_type = in_type +list_layer[0]*growth_rate self.trans1 = TransitionBlock(in_type, int(in_type/2), self.gspace) in_type = int(in_type/2) #2nd block self.block2 = DenseBlock(in_type, growth_rate, self.gspace, list_layer[1]) in_type = in_type +list_layer[1]*growth_rate self.trans2 = TransitionBlock(in_type, int(in_type/2), self.gspace) in_type = int(in_type/2) #3rd block self.block3 = DenseBlock(in_type, growth_rate, self.gspace, list_layer[2]) in_type = in_type +list_layer[2]*growth_rate self.trans3 = TransitionBlock(in_type, int(in_type/2), self.gspace) in_type = int(in_type/2) #4th block self.block4 = DenseBlock(in_type, growth_rate, self.gspace, list_layer[3]) in_type = in_type +list_layer[3]*growth_rate self.bn = enn.InnerBatchNorm(FIELD_TYPE["regular"](self.gspace, in_type, fixparams=False)) self.relu = enn.ReLU(FIELD_TYPE["regular"](self.gspace, in_type, fixparams=False),inplace=True) self.pool2 = torch.nn.AdaptiveAvgPool2d((1, 1)) self.classifier = torch.nn.Linear(in_type, nclasses)
def __init__(self, input_frames, output_frames, kernel_size, N): super(ResNet_Rot, self).__init__() r2_act = gspaces.Rot2dOnR2(N = N) # we use rho_1 representation since the input is velocity fields self.feat_type_in = nn.FieldType(r2_act, input_frames*[r2_act.irrep(1)]) # we use regular representation for middle layers self.feat_type_in_hid = nn.FieldType(r2_act, 16*[r2_act.regular_repr]) self.feat_type_hid_out = nn.FieldType(r2_act, 192*[r2_act.regular_repr]) self.feat_type_out = nn.FieldType(r2_act, output_frames*[r2_act.irrep(1)]) self.input_layer = nn.SequentialModule( nn.R2Conv(self.feat_type_in, self.feat_type_in_hid, kernel_size = kernel_size, padding = (kernel_size - 1)//2), nn.InnerBatchNorm(self.feat_type_in_hid), nn.ReLU(self.feat_type_in_hid) ) layers = [self.input_layer] layers += [rot_resblock(16, 32, kernel_size, N), rot_resblock(32, 32, kernel_size, N)] layers += [rot_resblock(32, 64, kernel_size, N), rot_resblock(64, 64, kernel_size, N)] layers += [rot_resblock(64, 128, kernel_size, N), rot_resblock(128, 128, kernel_size, N)] layers += [rot_resblock(128, 192, kernel_size, N), rot_resblock(192, 192, kernel_size, N)] layers += [nn.R2Conv(self.feat_type_hid_out, self.feat_type_out, kernel_size = kernel_size, padding = (kernel_size - 1)//2)] self.model = torch.nn.Sequential(*layers)
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)
def __init__(self, hidden_reps_ids, kernel_sizes, dim_cov_est, context_rep_ids=[1], N=4, flip=False, non_linearity=["NormReLU"], max_frequency=30): ''' Input: hidden_reps_ids - list: encoding the hidden fiber representation (see give_fib_reps_from_ids) kernel_sizes - list of ints - sizes of kernels for convolutional layers dim_cov_est - dimension of covariance estimation, either 1,2,3 or 4 context_rep_ids - list: gives the input fiber representation (see give_fib_reps_from_ids) non_linearity - list of strings - gives names of non-linearity to be used Either length 1 (then same non-linearity for all) or length is the number of layers (giving a custom non-linearity for every layer) N - int - gives the group order, -1 is infinite flip - Bool - indicates whether we have a flip in the rotation group (i.e.O(2) vs SO(2), D_N vs C_N) max_frequency - int - maximum irrep frequency to computed, only relevant if N=-1 ''' super(SteerDecoder, self).__init__() #Save the rotation group, if flip is true, then include all corresponding reflections: self.flip = flip self.max_frequency = max_frequency if self.flip: self.G_act = gspaces.FlipRot2dOnR2( N=N) if N != -1 else gspaces.FlipRot2dOnR2( N=N, maximum_frequency=self.max_frequency) #The output fiber representation is the identity: self.target_rep = self.G_act.irrep(1, 1) else: self.G_act = gspaces.Rot2dOnR2( N=N) if N != -1 else gspaces.Rot2dOnR2( N=N, maximum_frequency=self.max_frequency) #The output fiber representation is the identity: self.target_rep = self.G_act.irrep(1) #Save the N defining D_N or C_N (if N=-1 it is infinity): self.polygon_corners = N #Save the id's for the context representation and extract the context fiber representation: self.context_rep_ids = context_rep_ids self.context_rep = group.directsum( self.give_reps_from_ids(self.context_rep_ids)) #Save the parameters: self.kernel_sizes = kernel_sizes self.n_layers = len(hidden_reps_ids) + 2 self.hidden_reps_ids = hidden_reps_ids self.dim_cov_est = dim_cov_est #-----CREATE LIST OF NON-LINEARITIES---- if len(non_linearity) == 1: self.non_linearity = (self.n_layers - 2) * non_linearity elif len(non_linearity) != (self.n_layers - 2): sys.exit( "List of non-linearities invalid: must have either length 1 or n_layers-2" ) else: self.non_linearity = non_linearity #-----ENDE LIST OF NON-LINEARITIES---- #-----------CREATE DECODER----------------- ''' Create a list of layers based on the kernel sizes. Compute the padding such that the height h and width w of a tensor with shape (batch_size,n_channels,h,w) does not change while being passed through the decoder ''' #Create list of feature types: feat_types = self.give_feat_types() self.feature_emb = feat_types[0] self.feature_out = feat_types[-1] #Create layers list and append it: layers_list = [ G_CNN.R2Conv(feat_types[0], feat_types[1], kernel_size=kernel_sizes[0], padding=(kernel_sizes[0] - 1) // 2) ] for it in range(self.n_layers - 2): if self.non_linearity[it] == "ReLU": layers_list.append(G_CNN.ReLU(feat_types[it + 1], inplace=True)) elif self.non_linearity[it] == "NormReLU": layers_list.append(G_CNN.NormNonLinearity(feat_types[it + 1])) else: sys.exit("Unknown non-linearity.") layers_list.append( G_CNN.R2Conv(feat_types[it + 1], feat_types[it + 2], kernel_size=kernel_sizes[it], padding=(kernel_sizes[it] - 1) // 2)) #Create a steerable decoder out of the layers list: self.decoder = G_CNN.SequentialModule(*layers_list) #-----------END CREATE DECODER--------------- #-----------CONTROL INPUTS------------------ #Control that all kernel sizes are odd (otherwise output shape is not correct): if any([j % 2 - 1 for j in kernel_sizes]): sys.exit("All kernels need to have odd sizes") if len(kernel_sizes) != (self.n_layers - 1): sys.exit("Number of layers and number kernels do not match.") if len(self.non_linearity) != (self.n_layers - 2): sys.exit( "Number of layers and number of non-linearities do not match.")
def __init__(self, base = 'DNSteerableAGRadGalNet', attention_module='SelfAttention', attention_gates=3, attention_aggregation='ft', n_classes=2, attention_normalisation='sigmoid', quiet=True, number_rotations=8, imsize=150, kernel_size=3, group="D" ): super(DNSteerableAGRadGalNet, self).__init__() aggregation_mode = attention_aggregation normalisation = attention_normalisation AG = int(attention_gates) N = int(number_rotations) kernel_size = int(kernel_size) imsize = int(imsize) n_classes = int(n_classes) assert aggregation_mode in ['concat', 'mean', 'deep_sup', 'ft'], 'Aggregation mode not recognised. Valid inputs include concat, mean, deep_sup or ft.' assert normalisation in ['sigmoid','range_norm','std_mean_norm','tanh','softmax'], f'Nomralisation not implemented. Can be any of: sigmoid, range_norm, std_mean_norm, tanh, softmax' assert AG in [0,1,2,3], f'Number of Attention Gates applied (AG) must be an integer in range [0,3]. Currently AG={AG}' assert group.lower() in ["d","c"], f"group parameter must either be 'D' for DN, or 'C' for CN, steerable networks. (currently {group})." filters = [6,16,32,64,128] self.attention_out_sizes = [] self.ag = AG self.n_classes = n_classes self.filters = filters self.aggregation_mode = aggregation_mode # Setting up e2 if group.lower() == "d": self.r2_act = gspaces.FlipRot2dOnR2(N=int(number_rotations)) else: self.r2_act = gspaces.Rot2dOnR2(N=int(number_rotations)) in_type = e2nn.FieldType(self.r2_act, [self.r2_act.trivial_repr]) out_type = e2nn.FieldType(self.r2_act, 6*[self.r2_act.regular_repr]) self.in_type = in_type self.mask = e2nn.MaskModule(in_type, imsize, margin=0) self.conv1a = e2nn.R2Conv(in_type, out_type, kernel_size=kernel_size, padding=kernel_size//2, stride=1, bias=False); self.relu1a = e2nn.ReLU(out_type); self.bnorm1a= e2nn.InnerBatchNorm(out_type) self.conv1b = e2nn.R2Conv(out_type, out_type, kernel_size=kernel_size, padding=kernel_size//2, stride=1, bias=False); self.relu1b = e2nn.ReLU(out_type); self.bnorm1b= e2nn.InnerBatchNorm(out_type) self.conv1c = e2nn.R2Conv(out_type, out_type, kernel_size=kernel_size, padding=kernel_size//2, stride=1, bias=False); self.relu1c = e2nn.ReLU(out_type); self.bnorm1c= e2nn.InnerBatchNorm(out_type) self.mpool1 = e2nn.PointwiseMaxPool(out_type, kernel_size=(2,2), stride=2) self.gpool1 = e2nn.GroupPooling(out_type) in_type = out_type out_type = e2nn.FieldType(self.r2_act, 16*[self.r2_act.regular_repr]) self.conv2a = e2nn.R2Conv(in_type, out_type, kernel_size=kernel_size, padding=kernel_size//2, stride=1, bias=False); self.relu2a = e2nn.ReLU(out_type); self.bnorm2a= e2nn.InnerBatchNorm(out_type) self.conv2b = e2nn.R2Conv(out_type, out_type, kernel_size=kernel_size, padding=kernel_size//2, stride=1, bias=False); self.relu2b = e2nn.ReLU(out_type); self.bnorm2b= e2nn.InnerBatchNorm(out_type) self.conv2c = e2nn.R2Conv(out_type, out_type, kernel_size=kernel_size, padding=kernel_size//2, stride=1, bias=False); self.relu2c = e2nn.ReLU(out_type); self.bnorm2c= e2nn.InnerBatchNorm(out_type) self.mpool2 = e2nn.PointwiseMaxPool(out_type, kernel_size=(2,2), stride=2) self.gpool2 = e2nn.GroupPooling(out_type) in_type = out_type out_type = e2nn.FieldType(self.r2_act, 32*[self.r2_act.regular_repr]) self.conv3a = e2nn.R2Conv(in_type, out_type, kernel_size=kernel_size, padding=kernel_size//2, stride=1, bias=False); self.relu3a = e2nn.ReLU(out_type); self.bnorm3a= e2nn.InnerBatchNorm(out_type) self.conv3b = e2nn.R2Conv(out_type, out_type, kernel_size=kernel_size, padding=kernel_size//2, stride=1, bias=False); self.relu3b = e2nn.ReLU(out_type); self.bnorm3b= e2nn.InnerBatchNorm(out_type) self.conv3c = e2nn.R2Conv(out_type, out_type, kernel_size=kernel_size, padding=kernel_size//2, stride=1, bias=False); self.relu3c = e2nn.ReLU(out_type); self.bnorm3c= e2nn.InnerBatchNorm(out_type) self.mpool3 = e2nn.PointwiseMaxPool(out_type, kernel_size=(2,2), stride=2) self.gpool3 = e2nn.GroupPooling(out_type) in_type = out_type out_type = e2nn.FieldType(self.r2_act, 64*[self.r2_act.regular_repr]) self.conv4a = e2nn.R2Conv(in_type, out_type, kernel_size=kernel_size, padding=kernel_size//2, stride=1, bias=False); self.relu4a = e2nn.ReLU(out_type); self.bnorm4a= e2nn.InnerBatchNorm(out_type) self.conv4b = e2nn.R2Conv(out_type, out_type, kernel_size=kernel_size, padding=kernel_size//2, stride=1, bias=False); self.relu4b = e2nn.ReLU(out_type); self.bnorm4b= e2nn.InnerBatchNorm(out_type) self.mpool4 = e2nn.PointwiseMaxPool(out_type, kernel_size=(2,2), stride=2) self.gpool4 = e2nn.GroupPooling(out_type) self.flatten = nn.Flatten(1) self.dropout = nn.Dropout(p=0.5) if self.ag == 0: pass if self.ag >= 1: self.attention1 = GridAttentionBlock2D(in_channels=32, gating_channels=64, inter_channels=64, input_size=[imsize//4,imsize//4], normalisation=normalisation) if self.ag >= 2: self.attention2 = GridAttentionBlock2D(in_channels=16, gating_channels=64, inter_channels=64, input_size=[imsize//2,imsize//2], normalisation=normalisation) if self.ag >= 3: self.attention3 = GridAttentionBlock2D(in_channels=6, gating_channels=64, inter_channels=64, input_size=[imsize,imsize], normalisation=normalisation) self.fc1 = nn.Linear(16*5*5,256) #channel_size * width * height self.fc2 = nn.Linear(256,256) self.fc3 = nn.Linear(256, self.n_classes) self.dummy = nn.Parameter(torch.empty(0)) self.module_order = ['conv1a', 'relu1a', 'bnorm1a', #1->6 'conv1b', 'relu1b', 'bnorm1b', #6->6 'conv1c', 'relu1c', 'bnorm1c', #6->6 'mpool1', 'conv2a', 'relu2a', 'bnorm2a', #6->16 'conv2b', 'relu2b', 'bnorm2b', #16->16 'conv2c', 'relu2c', 'bnorm2c', #16->16 'mpool2', 'conv3a', 'relu3a', 'bnorm3a', #16->32 'conv3b', 'relu3b', 'bnorm3b', #32->32 'conv3c', 'relu3c', 'bnorm3c', #32->32 'mpool3', 'conv4a', 'relu4a', 'bnorm4a', #32->64 'conv4b', 'relu4b', 'bnorm4b', #64->64 'compatibility_score1', 'compatibility_score2'] ######################### # Aggreagation Strategies if self.ag != 0: self.attention_filter_sizes = [32, 16, 6] concat_length = 0 for i in range(self.ag): concat_length += self.attention_filter_sizes[i] if aggregation_mode == 'concat': self.classifier = nn.Linear(concat_length, self.n_classes) self.aggregate = self.aggregation_concat else: # Not able to initialise in a loop as the modules will not change device with remaining model. self.classifiers = nn.ModuleList() if self.ag>=1: self.classifiers.append(nn.Linear(self.attention_filter_sizes[0], self.n_classes)) if self.ag>=2: self.classifiers.append(nn.Linear(self.attention_filter_sizes[1], self.n_classes)) if self.ag>=3: self.classifiers.append(nn.Linear(self.attention_filter_sizes[2], self.n_classes)) if aggregation_mode == 'mean': self.aggregate = self.aggregation_sep elif aggregation_mode == 'deep_sup': self.classifier = nn.Linear(concat_length, self.n_classes) self.aggregate = self.aggregation_ds elif aggregation_mode == 'ft': self.classifier = nn.Linear(self.n_classes*self.ag, self.n_classes) self.aggregate = self.aggregation_ft else: raise NotImplementedError else: self.classifier = nn.Linear((150//16)**2*64, self.n_classes) self.aggregate = lambda x: self.classifier(self.flatten(x))
def __init__(self, in_type, out_type): super(Conv, self).__init__() self.conv = enn.SequentialModule( enn.R2Conv(in_type, out_type, kernel_size=3, stride=1, padding=0), enn.InnerBatchNorm(out_type), enn.ReLU(out_type))
def ennReLU(inplanes): in_type = FIELD_TYPE['regular'](gspace, inplanes) return enn.ReLU(in_type, inplace=True)
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, depth, widen_factor, dropout_rate, num_classes=100, N: int = 8, r: int = 1, f: bool = True, deltaorth: bool = False, fixparams: bool = True, initial_stride: int = 1, ): r""" Build and equivariant Wide ResNet. The parameter ``N`` controls rotation equivariance and the parameter ``f`` reflection equivariance. More precisely, ``N`` is the number of discrete rotations the model is initially equivariant to. ``N = 1`` means the model is only reflection equivariant from the beginning. ``f`` is a boolean flag specifying whether the model should be reflection equivariant or not. If it is ``False``, the model is not reflection equivariant. ``r`` is the restriction level: - ``0``: no restriction. The model is equivariant to ``N`` rotations from the input to the output - ``1``: restriction before the last block. The model is equivariant to ``N`` rotations before the last block (i.e. in the first 2 blocks). Then it is restricted to ``N/2`` rotations until the output. - ``2``: restriction after the first block. The model is equivariant to ``N`` rotations in the first block. Then it is restricted to ``N/2`` rotations until the output (i.e. in the last 3 blocks). - ``3``: restriction after the first and the second block. The model is equivariant to ``N`` rotations in the first block. It is restricted to ``N/2`` rotations before the second block and to ``1`` rotations before the last block. NOTICE: if restriction to ``N/2`` is performed, ``N`` needs to be even! """ super(Wide_ResNet, self).__init__() assert ((depth - 4) % 6 == 0), 'Wide-resnet depth should be 6n+4' n = int((depth - 4) / 6) k = widen_factor print(f'| Wide-Resnet {depth}x{k}') nStages = [16, 16 * k, 32 * k, 64 * k] self._fixparams = fixparams self._layer = 0 # number of discrete rotations to be equivariant to self._N = N # if the model is [F]lip equivariant self._f = f if self._f: if N != 1: self.gspace = gspaces.FlipRot2dOnR2(N) else: self.gspace = gspaces.Flip2dOnR2() else: if N != 1: self.gspace = gspaces.Rot2dOnR2(N) else: self.gspace = gspaces.TrivialOnR2() # level of [R]estriction: # r = 0: never do restriction, i.e. initial group (either DN or CN) preserved for the whole network # r = 1: restrict before the last block, i.e. initial group (either DN or CN) preserved for the first # 2 blocks, then restrict to N/2 rotations (either D{N/2} or C{N/2}) in the last block # r = 2: restrict after the first block, i.e. initial group (either DN or CN) preserved for the first # block, then restrict to N/2 rotations (either D{N/2} or C{N/2}) in the last 2 blocks # r = 3: restrict after each block. Initial group (either DN or CN) preserved for the first # block, then restrict to N/2 rotations (either D{N/2} or C{N/2}) in the second block and to 1 rotation # in the last one (D1 or C1) assert r in [0, 1, 2, 3] self._r = r # the input has 3 color channels (RGB). # Color channels are trivial fields and don't transform when the input is rotated or flipped r1 = enn.FieldType(self.gspace, [self.gspace.trivial_repr] * 3) # input field type of the model self.in_type = r1 # in the first layer we always scale up the output channels to allow for enough independent filters r2 = FIELD_TYPE["regular"](self.gspace, nStages[0], fixparams=True) # dummy attribute keeping track of the output field type of the last submodule built, i.e. the input field type of # the next submodule to build self._in_type = r2 self.conv1 = conv5x5(r1, r2) self.layer1 = self._wide_layer(WideBasic, nStages[1], n, dropout_rate, stride=initial_stride) if self._r >= 2: N_new = N//2 id = (0, N_new) if self._f else N_new self.restrict1 = self._restrict_layer(id) else: self.restrict1 = lambda x: x self.layer2 = self._wide_layer(WideBasic, nStages[2], n, dropout_rate, stride=2) if self._r == 3: id = (0, 1) if self._f else 1 self.restrict2 = self._restrict_layer(id) elif self._r == 1: N_new = N // 2 id = (0, N_new) if self._f else N_new self.restrict2 = self._restrict_layer(id) else: self.restrict2 = lambda x: x # last layer maps to a trivial (invariant) feature map self.layer3 = self._wide_layer(WideBasic, nStages[3], n, dropout_rate, stride=2, totrivial=True) self.bn = enn.InnerBatchNorm(self.layer3.out_type, momentum=0.9) self.relu = enn.ReLU(self.bn.out_type, inplace=True) self.linear = torch.nn.Linear(self.bn.out_type.size, num_classes) for name, module in self.named_modules(): if isinstance(module, enn.R2Conv): if deltaorth: init.deltaorthonormal_init(module.weights, module.basisexpansion) elif isinstance(module, torch.nn.BatchNorm2d): module.weight.data.fill_(1) module.bias.data.zero_() elif isinstance(module, torch.nn.Linear): module.bias.data.zero_() print("MODEL TOPOLOGY:") for i, (name, mod) in enumerate(self.named_modules()): print(f"\t{i} - {name}")
return x if __name__ == "__main__": r2_act = gspaces.Rot2dOnR2(N=8) # the input image is a scalar field, corresponding to the trivial representation in_type = nn.FieldType(r2_act, 3 * [r2_act.trivial_repr]) device = 'cuda' if torch.cuda.is_available() else 'cpu' dummy = torch.rand(64, 3, 32, 32).to(device) y = nn.GeometricTensor(dummy, in_type) # TEST 2 r2_act = gspaces.Rot2dOnR2(N=16) # 5 print(r2_act.trivial_repr) feat_type_in = nn.FieldType(r2_act, 3 * [r2_act.trivial_repr]) # 6 feat_type_out = nn.FieldType(r2_act, 10 * [r2_act.regular_repr]) # 7 # 8 conv = nn.R2Conv(feat_type_in, feat_type_out, kernel_size=5) # 9 relu = nn.ReLU(feat_type_out) # 10 # 11 x = torch.randn(16, 3, 32, 32) # 12 x = nn.GeometricTensor(x, feat_type_in) # 13 # 14 conv_x = conv(x) print(conv_x.size()) y = relu(conv_x) print(y.size())
def __init__(self, conv_func, group, in_channels): super(Backbone5x5, self).__init__() # the model is equivariant under rotations by 45 degrees, modelled by C8 # self.r2_act = gspaces.Rot2dOnR2(N=8) self.r2_act = group # the input image is a scalar field, corresponding to the trivial representation in_type = nn.FieldType(self.r2_act, in_channels * [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 if isinstance(in_type.gspace, e2cnn.gspaces.Rot2dOnR2): base = 8 elif isinstance(in_type.gspace, e2cnn.gspaces.FlipRot2dOnR2): base = 4 # 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, 16 * [self.r2_act.regular_repr]) self.add_module( 'block1', nn.SequentialModule( # nn.MaskModule(in_type, 29, margin=1), conv_func(in_type, out_type, kernel_size=5, 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 = out_type # the output type of the second convolution layer are 32 regular feature fields of C8 out_type = nn.FieldType(self.r2_act, 3 * base * [self.r2_act.regular_repr]) self.add_module( 'block2', nn.SequentialModule( conv_func(in_type, out_type, kernel_size=5, padding=2, bias=False), # nn.InnerBatchNorm(out_type), nn.ReLU(out_type, inplace=True))) self.add_module( 'pool1', nn.SequentialModule( nn.PointwiseMaxPool(out_type, kernel_size=3, stride=2))) # convolution 3 # the old output type is the input type to the next layer in_type = out_type # the output type of the third convolution layer are 32 regular feature fields of C8 out_type = nn.FieldType(self.r2_act, 4 * base * [self.r2_act.regular_repr]) self.add_module( 'block3', nn.SequentialModule( conv_func(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 = out_type # the output type of the fourth convolution layer are 64 regular feature fields of C8 out_type = nn.FieldType(self.r2_act, 6 * base * [self.r2_act.regular_repr]) self.add_module( 'block4', nn.SequentialModule( conv_func(in_type, out_type, kernel_size=5, padding=2, bias=False), # nn.InnerBatchNorm(out_type), nn.ReLU(out_type, inplace=True))) self.add_module( 'pool2', nn.SequentialModule( nn.PointwiseMaxPool(out_type, kernel_size=3, stride=2))) # convolution 5 # the old output type is the input type to the next layer in_type = out_type # the output type of the fifth convolution layer are 64 regular feature fields of C8 out_type = nn.FieldType(self.r2_act, 8 * base * [self.r2_act.regular_repr]) self.add_module( 'block5', nn.SequentialModule( conv_func(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 = out_type # the output type of the sixth convolution layer are 64 regular feature fields of C8 out_type = nn.FieldType(self.r2_act, 12 * base * [self.r2_act.regular_repr]) self.add_module( 'block6', nn.SequentialModule( conv_func(in_type, out_type, kernel_size=5, padding=1, bias=False), # nn.InnerBatchNorm(out_type), nn.ReLU(out_type, inplace=True))) self.add_module( 'pool3', nn.PointwiseMaxPool(out_type, kernel_size=3, stride=1, padding=0)) self.out_type = out_type