def make_layers_chip0(args, use_bn=True):
mask_bits = spec.specs[args.chip]['resnet50']
return nn.Sequential(
basic_conv_block(3,
64,
downsample_mode="MAXPOOL",
block_size=2,
use_bn=use_bn,
quant_params=args,
mask_bit=mask_bits[0]),
basic_conv_block(64,
64,
downsample_mode="MAXPOOL",
block_size=2,
use_bn=use_bn,
quant_params=args,
mask_bit=mask_bits[1]),
residual_block(64,
64,
downsample_mode="MAXPOOL",
block_size=3,
use_bn=use_bn,
quant_params=args,
mask_bit=mask_bits[2]),
basic_conv_block(64,
128,
downsample_mode=None,
block_size=2,
use_bn=use_bn,
quant_params=args,
mask_bit=mask_bits[3]),
residual_block(128,
128,
downsample_mode="MAXPOOL",
block_size=3,
use_bn=use_bn,
quant_params=args,
mask_bit=mask_bits[3]),
basic_conv_block(128,
256,
downsample_mode=None,
block_size=2,
use_bn=use_bn,
quant_params=args,
mask_bit=mask_bits[4]),
residual_block(256,
256,
downsample_mode=None,
block_size=1,
use_bn=use_bn,
quant_params=args,
mask_bit=mask_bits[4]),
residual_block(256,
256,
downsample_mode=None,
block_size=6,
use_bn=use_bn,
quant_params=args,
mask_bit=mask_bits[5]))
def make_hyper_encoder_chip0(args):
mask_bits = spec.specs[args.chip]['hyper_encoder']
return nn.Sequential(
residual_block(256, 256, downsample_mode=None,
block_size=3, use_bn=False, quant_params=args, mask_bit=mask_bits[0]),
basic_conv_block(256, 256, downsample_mode="STRIDE2",
block_size=2, use_bn=False, quant_params=args, mask_bit=mask_bits[0]),
residual_block(256, 256, downsample_mode=None,
block_size=3, use_bn=False, quant_params=args, mask_bit=mask_bits[1]),
)
def make_layers(args):
mask_bits = spec.specs[args.chip]['tenbit']
return nn.Sequential(
basic_conv_block(3,
64,
downsample_mode="MAXPOOL",
block_size=2,
quant_params=args,
mask_bit=mask_bits[0]),
basic_conv_block(64,
64,
downsample_mode="MAXPOOL",
block_size=2,
quant_params=args,
mask_bit=mask_bits[1]),
residual_block(64,
64,
downsample_mode="MAXPOOL",
block_size=2,
quant_params=args,
mask_bit=mask_bits[2]),
basic_conv_block(64,
128,
downsample_mode=None,
block_size=2,
quant_params=args,
mask_bit=mask_bits[3]),
residual_block(128,
128,
downsample_mode="MAXPOOL",
block_size=1,
quant_params=args,
mask_bit=mask_bits[3]),
basic_conv_block(128,
256,
downsample_mode=None,
block_size=2,
quant_params=args,
mask_bit=mask_bits[4]),
residual_block(256,
256,
downsample_mode="MAXPOOL",
block_size=1,
quant_params=args,
mask_bit=mask_bits[4]),
basic_conv_block(256,
512,
downsample_mode=None,
block_size=1,
quant_params=args,
mask_bit=mask_bits[5]),
)
def make_hyper_decoder_host0(args):
return nn.Sequential(
residual_block(256, 256, downsample_mode=None,
block_size=3, use_bn=False, quant_params=args, mask_bit=host_mask_bit),
basic_conv_block(256, 256, downsample_mode=None,
block_size=1, use_bn=False, quant_params=args, mask_bit=host_mask_bit),
deconv_block(256, 256, upsampling_mode=_UPSAMPLING_MODE,
block_size=1, use_bn=False, quant_params=args, mask_bit=host_mask_bit)
)
def make_main_decoder_chip0(args):
mask_bits = spec.specs[args.chip]['main_decoder']
return nn.Sequential(
basic_conv_block(256, 256, downsample_mode=None,
block_size=1, use_bn=False, quant_params=args, mask_bit=mask_bits[0]),
deconv_block(256, 256, upsampling_mode=_UPSAMPLING_MODE,
block_size=1, use_bn=False, quant_params=args, mask_bit=mask_bits[0]),
residual_block(256, 256, downsample_mode=None,
block_size=3, use_bn=False, quant_params=args, mask_bit=mask_bits[1]),
basic_conv_block(256, 256, downsample_mode=None,
block_size=1, use_bn=False, quant_params=args, mask_bit=mask_bits[1]),
deconv_block(256, 256, upsampling_mode=_UPSAMPLING_MODE,
block_size=1, use_bn=False, quant_params=args, mask_bit=mask_bits[1]),
residual_block(256, 256, downsample_mode=None,
block_size=3, use_bn=False, quant_params=args, mask_bit=mask_bits[2]),
basic_conv_block(256, 256, downsample_mode=None,
block_size=3, use_bn=False, quant_params=args, mask_bit=mask_bits[2]),
deconv_block(256, 256, upsampling_mode=_UPSAMPLING_MODE,
block_size=1, use_bn=False, quant_params=args, mask_bit=mask_bits[2])
)
def make_layers_chip1(args, use_bn=True):
mask_bits = spec.specs[args.chip]['resnet50']
return nn.Sequential(
basic_conv_block(256,
512,
downsample_mode="STRIDE2",
block_size=1,
use_bn=use_bn,
quant_params=args,
mask_bit=mask_bits[6]),
basic_conv_block(512,
512,
downsample_mode=None,
block_size=1,
use_bn=use_bn,
quant_params=args,
mask_bit=mask_bits[7]),
residual_block(512,
512,
downsample_mode=None,
block_size=2,
use_bn=use_bn,
quant_params=args,
mask_bit=mask_bits[8]))
def make_hyper_decoder_chip0(args):
mask_bits = spec.specs[args.chip]['hyper_decoder']
return nn.Sequential(
residual_block(256, 256, downsample_mode=None,
block_size=3, use_bn=False, quant_params=args, mask_bit=mask_bits[0])
)