1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69
| import torch import torch.nn as nn
class MixedPrecisionDMS(nn.Module): """混合精度DMS模型""" def __init__(self): super().__init__() self.backbone = nn.Sequential( nn.Conv2d(3, 64, 7, stride=2, padding=3), nn.BatchNorm2d(64), nn.ReLU(), QuantizedConv2d(64, 128, 3, padding=1, precision="INT8"), nn.BatchNorm2d(128), nn.ReLU(), QuantizedConv2d(128, 256, 3, padding=1, precision="INT4"), nn.BatchNorm2d(256), nn.ReLU() ) self.head = nn.Linear(256, 7) def forward(self, x): x = self.backbone(x) x = self.head(x) return x
class QuantizedConv2d(nn.Module): """量化卷积层""" def __init__(self, in_channels, out_channels, kernel_size, padding=0, precision="INT8"): super().__init__() self.precision = precision self.bit_width = { "INT2": 2, "INT4": 4, "INT8": 8, "INT16": 16, "FP8": 8, "FP16": 16 }[precision] self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, padding=padding) def forward(self, x): if self.precision.startswith("INT"): scale = x.abs().max() / (2 ** (self.bit_width - 1) - 1) x = torch.round(x / scale) * scale elif self.precision == "FP8": x = x.to(torch.float8_e4m3fn) return self.conv(x)
|