AlexNet模型使用mindspore框架重新定义

tests/model/AlexNet.py

@@ -0,0 +1,77 @@
+import mindspore
+import mindspore.nn as nn
+import mindspore.ops.operations as P
+
+
+class AlexNet(nn.Cell):
+    def __init__(self, input_channels, output_num, input_size):
+        super().__init__()
+
+        self.features = nn.SequentialCell([
+            nn.Conv2d(in_channels=input_channels, out_channels=64, kernel_size=3, stride=2, pad_mode='pad', padding=1,
+                      has_bias=True),
+            nn.BatchNorm2d(num_features=64, momentum=0.9),  # 批量归一化层
+            nn.MaxPool2d(kernel_size=2, stride=2, pad_mode='valid'),
+            nn.ReLU(),
+
+            nn.Conv2d(in_channels=64, out_channels=192, kernel_size=3, pad_mode='pad', padding=1, has_bias=True),
+            nn.BatchNorm2d(num_features=192, momentum=0.9),  # 批量归一化层
+            nn.MaxPool2d(kernel_size=2, stride=2, pad_mode='valid'),
+            nn.ReLU(),
+
+            nn.Conv2d(in_channels=192, out_channels=384, kernel_size=3, pad_mode='pad', padding=1, has_bias=True),
+            nn.BatchNorm2d(num_features=384, momentum=0.9),  # 批量归一化层
+            nn.ReLU(),
+
+            nn.Conv2d(in_channels=384, out_channels=256, kernel_size=3, pad_mode='pad', padding=1, has_bias=True),
+            nn.BatchNorm2d(num_features=256, momentum=0.9),  # 批量归一化层
+            nn.ReLU(),
+
+            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, pad_mode='pad', padding=1, has_bias=True),
+            nn.BatchNorm2d(num_features=256, momentum=0.9),  # 批量归一化层
+            nn.MaxPool2d(kernel_size=2, stride=2, pad_mode='valid'),
+            nn.ReLU(),
+        ])
+
+        self.input_size = input_size
+        self._init_classifier(output_num)
+
+    def _init_classifier(self, output_num):
+        # Forward a dummy input through the feature extractor part of the network
+        dummy_input = mindspore.ops.zeros((1, 3, self.input_size, self.input_size))
+        features_size = self.features(dummy_input).numel()
+
+        self.classifier = nn.SequentialCell([
+            nn.Dropout(p=0.5),
+            nn.Dense(in_channels=features_size, out_channels=1000),
+            nn.ReLU(),
+
+            nn.Dropout(p=0.5),
+            nn.Dense(in_channels=1000, out_channels=256),
+            nn.ReLU(),
+
+            nn.Dense(in_channels=256, out_channels=output_num)
+        ])
+
+    def construct(self, x):
+        x = self.features(x)
+        x = P.Reshape()(x, (P.Shape()(x)[0], -1,))
+        x = self.classifier(x)
+        return x
+
+
+if __name__ == '__main__':
+    import argparse
+
+    parser = argparse.ArgumentParser(description='AlexNet Implementation')
+    parser.add_argument('--input_channels', default=3, type=int)
+    parser.add_argument('--output_num', default=10, type=int)
+    parser.add_argument('--input_size', default=32, type=int)
+    args = parser.parse_args()
+
+    model = AlexNet(args.input_channels, args.output_num, args.input_size)
+    tensor = mindspore.ops.rand((1, args.input_channels, args.input_size, args.input_size))
+    pred = model(tensor)
+
+    print(model)
+    print("Predictions shape:", pred.shape)