| 1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677 |
- 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)
|
