| @@ -142,33 +142,33 @@ class QuantizationAwareTraining(Quantizer): | |||
| Examples: | |||
| >>> class LeNet5(nn.Cell): | |||
| >>> def __init__(self, num_class=10, channel=1): | |||
| >>> super(LeNet5, self).__init__() | |||
| >>> self.type = "fusion" | |||
| >>> self.num_class = num_class | |||
| >>> | |||
| >>> # change `nn.Conv2d` to `nn.Conv2dBnAct` | |||
| >>> self.conv1 = nn.Conv2dBnAct(channel, 6, 5, pad_mode='valid', activation='relu') | |||
| >>> self.conv2 = nn.Conv2dBnAct(6, 16, 5, pad_mode='valid', activation='relu') | |||
| >>> # change `nn.Dense` to `nn.DenseBnAct` | |||
| >>> self.fc1 = nn.DenseBnAct(16 * 5 * 5, 120, activation='relu') | |||
| >>> self.fc2 = nn.DenseBnAct(120, 84, activation='relu') | |||
| >>> self.fc3 = nn.DenseBnAct(84, self.num_class) | |||
| >>> | |||
| >>> self.max_pool2d = nn.MaxPool2d(kernel_size=2, stride=2) | |||
| >>> self.flatten = nn.Flatten() | |||
| >>> | |||
| >>> def construct(self, x): | |||
| >>> x = self.conv1(x) | |||
| >>> x = self.max_pool2d(x) | |||
| >>> x = self.conv2(x) | |||
| >>> x = self.max_pool2d(x) | |||
| >>> x = self.flatten(x) | |||
| >>> x = self.fc1(x) | |||
| >>> x = self.fc2(x) | |||
| >>> x = self.fc3(x) | |||
| >>> return x | |||
| >>> | |||
| ... def __init__(self, num_class=10, channel=1): | |||
| ... super(LeNet5, self).__init__() | |||
| ... self.type = "fusion" | |||
| ... self.num_class = num_class | |||
| ... | |||
| ... # change `nn.Conv2d` to `nn.Conv2dBnAct` | |||
| ... self.conv1 = nn.Conv2dBnAct(channel, 6, 5, pad_mode='valid', activation='relu') | |||
| ... self.conv2 = nn.Conv2dBnAct(6, 16, 5, pad_mode='valid', activation='relu') | |||
| ... # change `nn.Dense` to `nn.DenseBnAct` | |||
| ... self.fc1 = nn.DenseBnAct(16 * 5 * 5, 120, activation='relu') | |||
| ... self.fc2 = nn.DenseBnAct(120, 84, activation='relu') | |||
| ... self.fc3 = nn.DenseBnAct(84, self.num_class) | |||
| ... | |||
| ... self.max_pool2d = nn.MaxPool2d(kernel_size=2, stride=2) | |||
| ... self.flatten = nn.Flatten() | |||
| ... | |||
| ... def construct(self, x): | |||
| ... x = self.conv1(x) | |||
| ... x = self.max_pool2d(x) | |||
| ... x = self.conv2(x) | |||
| ... x = self.max_pool2d(x) | |||
| ... x = self.flatten(x) | |||
| ... x = self.fc1(x) | |||
| ... x = self.fc2(x) | |||
| ... x = self.fc3(x) | |||
| ... return x | |||
| ... | |||
| >>> net = LeNet5() | |||
| >>> quantizer = QuantizationAwareTraining(bn_fold=False, per_channel=[True, False], symmetric=[True, False]) | |||
| >>> net_qat = quantizer.quantize(net) | |||