| import torch |
| import torch.nn as nn |
| |
| |
| class GeneralQuantModule(torch.nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.embedding = torch.ao.nn.quantized.Embedding( |
| num_embeddings=10, embedding_dim=12 |
| ) |
| self.embedding_input = torch.tensor([9, 6, 5, 7, 8, 8, 9, 2, 8]) |
| self.func = torch.ao.nn.quantized.QFunctional() |
| self.conv1 = torch.ao.nn.quantized.ConvTranspose1d(16, 33, 3, stride=2) |
| self.conv2 = torch.ao.nn.quantized.ConvTranspose2d(16, 33, 3, stride=2) |
| self.conv3 = torch.ao.nn.quantized.ConvTranspose3d(16, 33, 3, stride=2) |
| |
| def forward(self): |
| a = torch.quantize_per_tensor(torch.tensor([3.0]), 1.0, 0, torch.qint32) |
| b = torch.quantize_per_tensor(torch.tensor(4.0), 1.0, 0, torch.qint32) |
| c = torch.quantize_per_tensor( |
| torch.tensor([3.0]), torch.tensor(1.0), torch.tensor(0), torch.qint32 |
| ) |
| input1 = torch.randn(1, 16, 4) |
| input2 = torch.randn(1, 16, 4, 4) |
| input3 = torch.randn(1, 16, 4, 4, 4) |
| return len( |
| self.func.add(a, b), |
| self.func.cat((a, a), 0), |
| self.func.mul(a, b), |
| self.func.add_relu(a, b), |
| self.func.add_scalar(a, b), |
| self.func.mul_scalar(a, b), |
| self.embedding(self.embedding_input), |
| self.conv1( |
| torch.quantize_per_tensor( |
| input1, scale=1.0, zero_point=0, dtype=torch.quint8 |
| ) |
| ), |
| self.conv2( |
| torch.quantize_per_tensor( |
| input2, scale=1.0, zero_point=0, dtype=torch.quint8 |
| ) |
| ), |
| c, |
| # self.conv3(torch.quantize_per_tensor(input3, scale=1.0, zero_point=0, dtype=torch.quint8)), # failed on iOS |
| ) |
| |
| |
| class DynamicQuantModule: |
| def __init__(self): |
| super().__init__() |
| self.module = self.M() |
| |
| def getModule(self): |
| return torch.ao.quantization.quantize_dynamic(self.module, dtype=torch.qint8) |
| |
| class M(torch.nn.Module): |
| def __init__(self): |
| super(DynamicQuantModule.M, self).__init__() |
| self.rnn = nn.RNN(4, 8, 2) |
| self.rnncell = nn.RNNCell(4, 8) |
| self.gru = nn.GRU(4, 8, 2) |
| self.grucell = nn.GRUCell(4, 8) |
| self.lstm = nn.LSTM(4, 8, 2) |
| self.lstmcell = nn.LSTMCell(4, 8) |
| self.linears = nn.ModuleList( |
| [ |
| nn.Identity(54), |
| nn.Linear(20, 20), |
| nn.Bilinear(20, 20, 40), |
| ] |
| ) |
| self.transformers = nn.ModuleList( |
| [ |
| nn.Transformer( |
| d_model=2, nhead=2, num_encoder_layers=1, num_decoder_layers=1 |
| ), |
| nn.TransformerEncoder( |
| nn.TransformerEncoderLayer(d_model=2, nhead=2), num_layers=1 |
| ), |
| nn.TransformerDecoder( |
| nn.TransformerDecoderLayer(d_model=2, nhead=2), num_layers=1 |
| ), |
| ] |
| ) |
| # self.a = torch.nn.utils.rnn.pad_sequence([torch.tensor([1,2,3]), torch.tensor([3,4])], batch_first=True) |
| |
| def forward(self): |
| input = torch.randn(5, 3, 4) |
| h = torch.randn(2, 3, 8) |
| c = torch.randn(2, 3, 8) |
| linear_input = torch.randn(32, 20) |
| trans_input = torch.randn(1, 16, 2) |
| tgt = torch.rand(1, 16, 2) |
| |
| return len(( |
| self.rnn(input, h), |
| self.rnncell(input[0], h[0]), |
| self.gru(input, h), |
| self.grucell(input[0], h[0]), |
| self.lstm(input, (h, c)), |
| # self.lstm(torch.nn.utils.rnn.pack_padded_sequence(self.a, lengths=torch.tensor([3,2,1])), (h, c)), |
| self.lstmcell(input[0], (h[0], c[0])), |
| self.transformers[0](trans_input, tgt), |
| self.transformers[1](trans_input), |
| self.transformers[2](trans_input, tgt), |
| self.linears[0](linear_input), |
| self.linears[1](linear_input), |
| self.linears[2](linear_input, linear_input), |
| )) |
| |
| |
| class StaticQuantModule: |
| def getModule(self): |
| model_fp32 = self.M() |
| model_fp32.eval() |
| model_fp32.qconfig = torch.ao.quantization.get_default_qconfig("qnnpack") |
| model_fp32_prepared = torch.ao.quantization.prepare(model_fp32) |
| model_int8 = torch.ao.quantization.convert(model_fp32_prepared) |
| return model_int8 |
| |
| class M(torch.nn.Module): |
| def __init__(self): |
| super(StaticQuantModule.M, self).__init__() |
| self.quant = torch.ao.quantization.QuantStub() |
| self.input1d = torch.randn(4, 2, 2) |
| self.input2d = torch.randn((4, 2, 4, 4)) |
| self.input3d = torch.randn(4, 2, 2, 4, 4) |
| self.linear_input = torch.randn(32, 20) |
| |
| self.layer1 = nn.Sequential( |
| nn.Conv1d(2, 2, 1), nn.InstanceNorm1d(1), nn.Hardswish() |
| ) |
| self.layer2 = nn.Sequential( |
| nn.Conv2d(2, 2, 1), |
| nn.BatchNorm2d(2), |
| nn.InstanceNorm2d(1), |
| nn.LeakyReLU(), |
| ) |
| self.layer3 = nn.Sequential( |
| nn.Conv3d(2, 2, 1), nn.BatchNorm3d(2), nn.InstanceNorm3d(1), nn.ReLU() |
| ) |
| self.layer4 = nn.Sequential(nn.Linear(4, 3)) |
| self.dequant = torch.ao.quantization.DeQuantStub() |
| |
| def forward(self): |
| x = self.quant(self.input1d) |
| x = self.layer1(x) |
| x = self.dequant(x) |
| |
| y = self.input2d |
| y = self.quant(y) |
| y = self.layer2(y) |
| y = self.layer4(y) |
| y = self.dequant(y) |
| |
| z = self.quant(self.input3d) |
| z = self.layer3(z) |
| z = self.dequant(z) |
| |
| return (x, y, z) |
| |
| |
| class FusedQuantModule: |
| def getModule(self): |
| model_fp32 = self.M() |
| model_fp32.eval() |
| model_fp32.qconfig = torch.ao.quantization.get_default_qconfig("qnnpack") |
| model_fp32_fused = torch.ao.quantization.fuse_modules( |
| model_fp32, |
| [ |
| ["conv1d", "relu1"], |
| ["conv2d", "relu2"], |
| ["conv3d", "relu3"], |
| ["linear", "relu4"], |
| ], |
| ) |
| model_fp32_prepared = torch.ao.quantization.prepare(model_fp32_fused) |
| model_int8 = torch.ao.quantization.convert(model_fp32_prepared) |
| return model_int8 |
| |
| class M(torch.nn.Module): |
| def __init__(self): |
| super(FusedQuantModule.M, self).__init__() |
| self.quant = torch.ao.quantization.QuantStub() |
| self.input1d = torch.randn(4, 2, 2) |
| self.input2d = torch.randn((4, 2, 4, 4)) |
| self.input3d = torch.randn(4, 2, 2, 4, 4) |
| self.conv1d = nn.Conv1d(2, 2, 1) |
| self.conv2d = nn.Conv2d(2, 2, 1) |
| self.conv3d = nn.Conv3d(2, 2, 1) |
| self.linear = nn.Linear(4, 2) |
| self.relu1 = nn.ReLU() |
| self.relu2 = nn.ReLU() |
| self.relu3 = nn.ReLU() |
| self.relu4 = nn.ReLU() |
| self.dequant = torch.ao.quantization.DeQuantStub() |
| |
| def forward(self): |
| x = self.input1d |
| y = self.input2d |
| z = self.input3d |
| |
| x = self.quant(x) |
| x = self.conv1d(x) |
| x = self.relu1(x) |
| x = self.dequant(x) |
| |
| y = self.quant(y) |
| y = self.conv2d(y) |
| y = self.relu2(y) |
| y = self.dequant(y) |
| |
| z = self.quant(z) |
| z = self.conv3d(z) |
| z = self.relu3(z) |
| z = self.linear(z) |
| z = self.relu4(z) |
| z = self.dequant(z) |
| |
| return (x, y, z) |
