def __init__(self,
enet_type,
out_dim,
pretrained=True,
freeze_cnn=False,
load_model=False,
pretrain_cnn=False,
pretrain_file=''):
super(Effnet, self).__init__()
self.enet = geffnet.create_model(enet_type, pretrained=pretrained)
in_ch = self.enet.classifier.in_features
self.enet.classifier = nn.Identity()
if pretrain_cnn:
pretrained_dict = torch.load(pretrain_file)
model_dict = self.state_dict()
pretrained_dict = {
k: v
for k, v in pretrained_dict.items() if k in model_dict
}
model_dict.update(pretrained_dict)
self.load_state_dict(model_dict)
self.dropouts = nn.ModuleList([nn.Dropout(0.5) for _ in range(5)])
self.myfc = nn.Linear(in_ch, out_dim)
if freeze_cnn:
for p in self.enet.parameters():
p.requires_grad = False
def __init__(self,
enet_type,
out_dim,
n_meta_features=0,
n_meta_dim=[512, 128],
pretrained=False):
super(Effnet_MMC, self).__init__()
self.n_meta_features = n_meta_features
# efficient net 모델
self.enet = geffnet.create_model(enet_type, pretrained=pretrained)
self.dropouts = nn.ModuleList([nn.Dropout(0.5) for _ in range(5)])
in_ch = self.enet.classifier.in_features
if n_meta_features > 0:
# meta 데이터를 사용한 경우, 짤막한 2레이어 분류기 추가
self.meta = nn.Sequential(
nn.Linear(n_meta_features, n_meta_dim[0]),
nn.BatchNorm1d(n_meta_dim[0]),
# swish activation function
Swish_Module(),
nn.Dropout(p=0.3),
nn.Linear(n_meta_dim[0], n_meta_dim[1]),
nn.BatchNorm1d(n_meta_dim[1]),
Swish_Module(),
)
in_ch += n_meta_dim[1]
self.myfc = nn.Linear(in_ch, out_dim)
self.enet.classifier = nn.Identity()
def __init__(self, enet_type, out_dim):
super(enet_arcface_FINAL, self).__init__()
self.enet = geffnet.create_model(enet_type.replace("-", "_"), pretrained=None)
self.feat = nn.Linear(self.enet.classifier.in_features, 512)
self.swish = Swish_module()
self.metric_classify = ArcMarginProduct_subcenter(512, out_dim)
self.enet.classifier = nn.Identity()
def create(self) -> Module:
return geffnet.create_model(
self.model_config.model_name,
drop_rate=self.model_config.drop_rate,
drop_connect_rate=self.model_config.drop_connect_rate,
pretrained=self.model_config.pretrained,
num_classes=4)
def __init__(self,
enet_type,
out_dim,
n_meta_features=0,
n_meta_dim=[512, 128],
pretrained=False):
super(Effnet_Melanoma, self).__init__()
self.n_meta_features = n_meta_features
self.enet = geffnet.create_model(
enet_type, pretrained=pretrained) # ! make EfficientNet
self.dropouts = nn.ModuleList([nn.Dropout(0.5) for _ in range(5)])
in_ch = self.enet.classifier.in_features
if n_meta_features > 0:
self.meta = nn.Sequential(
nn.Linear(n_meta_features, n_meta_dim[0]),
nn.BatchNorm1d(n_meta_dim[0]),
Swish_Module(),
nn.Dropout(p=0.3),
nn.Linear(n_meta_dim[0], n_meta_dim[1]),
nn.BatchNorm1d(n_meta_dim[1]),
Swish_Module(),
)
in_ch += n_meta_dim[1]
self.myfc = nn.Linear(in_ch, out_dim) # ! simple classifier
self.enet.classifier = nn.Identity() # ! pass through, no update
def run_single_test(model_name, input_size):
input = torch.randn(input_size)
# create model from geffnet
model = geffnet.create_model(model_name, pretrained=True)
model.eval()
# print(model)
with torch.no_grad():
for i in range(nwarmups):
output = model(input)
t1 = time()
with torch.autograd.profiler.profile(enabled=args.profile) as prof:
for i in range(niters):
output = model(input)
if args.profile:
print("\n\n{} profiling result:".format(model_name))
print(prof.key_averages().table(sort_by="self_cpu_time_total",
row_limit=20))
t2 = time()
ttime = (t2 - t1) / niters * 1000
print('Model: {}; Input size: {}; time: {:.3f} ms'.format(
model_name, input_size, ttime))
def __init__(self,
backbone,
rpn_head,
bbox_roi_extractor,
bbox_head,
mask_roi_extractor,
mask_head,
semantic_head,
train_cfg,
test_cfg,
neck=None,
shared_head=None,
pretrained=None):
assert backbone is not None
assert rpn_head is not None
assert bbox_roi_extractor is not None
assert bbox_head is not None
assert mask_roi_extractor is not None
assert mask_head is not None
assert semantic_head is not None
super(EfficientPS, self).__init__()
self.eff_backbone_flag = False if 'efficient' not in backbone[
'type'] else True
if self.eff_backbone_flag == False:
self.backbone = builder.build_backbone(backbone)
else:
self.backbone = geffnet.create_model(
backbone['type'],
pretrained=True if pretrained is not None else False,
se=False,
act_layer=backbone['act_cfg']['type'],
norm_layer=norm_cfg[backbone['norm_cfg']['type']][1])
self.neck = builder.build_neck(neck)
if shared_head is not None:
self.shared_head = builder.build_shared_head(shared_head)
self.rpn_head = builder.build_head(rpn_head)
self.bbox_roi_extractor = builder.build_roi_extractor(
bbox_roi_extractor)
self.bbox_head = builder.build_head(bbox_head)
self.mask_roi_extractor = builder.build_roi_extractor(
mask_roi_extractor)
self.share_roi_extractor = True
self.mask_head = builder.build_head(mask_head)
self.semantic_head = builder.build_head(semantic_head)
self.train_cfg = train_cfg
self.test_cfg = test_cfg
self.num_classes = semantic_head['num_classes']
self.num_stuff = self.num_classes - bbox_head['num_classes'] + 1
self.init_weights(pretrained=pretrained)
def main():
args = parser.parse_args()
if not args.checkpoint:
args.pretrained = True
else:
args.pretrained = False
# create model
geffnet.config.set_exportable(True)
print("==> Creating PyTorch {} model".format(args.model))
model = geffnet.create_model(args.model,
num_classes=args.num_classes,
in_chans=3,
pretrained=args.pretrained,
checkpoint_path=args.checkpoint)
model.eval()
x = torch.randn((1, 3, args.img_size or 224, args.img_size or 224),
requires_grad=True)
model(x) # run model once before export trace
print("==> Exporting model to ONNX format at '{}'".format(args.output))
input_names = ["input0"]
output_names = ["output0"]
optional_args = dict(keep_initializers_as_inputs=True
) # pytorch 1.3 needs this for export to succeed
try:
torch_out = torch.onnx._export(model,
x,
args.output,
export_params=True,
verbose=False,
input_names=input_names,
output_names=output_names,
**optional_args)
except TypeError:
# fallback to no keep_initializers arg for pytorch < 1.3
torch_out = torch.onnx._export(model,
x,
args.output,
export_params=True,
verbose=False,
input_names=input_names,
output_names=output_names)
print("==> Loading and checking exported model from '{}'".format(
args.output))
onnx_model = onnx.load(args.output)
onnx.checker.check_model(onnx_model) # assuming throw on error
print("==> Passed")
print("==> Loading model into Caffe2 backend and comparing forward pass.".
format(args.output))
caffe2_backend = onnx_caffe2.prepare(onnx_model)
B = {onnx_model.graph.input[0].name: x.data.numpy()}
c2_out = caffe2_backend.run(B)[0]
np.testing.assert_almost_equal(torch_out.data.numpy(), c2_out, decimal=5)
print("==> Passed")
def load_model(model_type, state_dict=None, for_train=False):
if model_type == 'res18':
netD = torchvision.models.resnet18(pretrained=False)
feature_size = netD.fc.in_features
netD.fc = nn.Linear(feature_size, 2)
elif model_type == 'densenet':
netD = torchvision.models.densenet121(pretrained=False)
netD.classifier = nn.Linear(1024, 2)
elif model_type == 'mobilenet':
netD = geffnet.create_model('mobilenetv3_rw', pretrained=False)
netD.classifier = nn.Linear(in_features=1280, out_features=2, bias=True)
if state_dict:
print('load state_dict')
netD.load_state_dict(torch.load(state_dict))
if for_train:
for params in netD.parameters():
requires_grad = True
netD.train()
else:
for params in netD.parameters():
requires_grad = False
netD.eval()
netD.cuda()
return netD
def __init__(self, enet_type, out_dim):
super(Effnet_Landmark, self).__init__()
self.enet = geffnet.create_model(enet_type.replace('-', '_'),
pretrained=True)
self.feat = nn.Linear(self.enet.classifier.in_features, 512)
self.swish = Swish_module()
self.metric_classify = ArcMarginProduct_subcenter(512, out_dim)
self.enet.classifier = nn.Identity()
def __init__(self, backbone, out_dim, n_meta_features=0, load_pretrained=True):
super(enetv2, self).__init__()
self.n_meta_features = n_meta_features
self.enet = geffnet.create_model(backbone, pretrained=True)
self.dropout = nn.Dropout(0.5)
in_ch = self.enet.classifier.in_features
self.myfc = nn.Linear(in_ch, out_dim)
self.enet.classifier = nn.Identity()
def mobilenetv2_bak(pretrained=True, **kwargs):
model = geffnet.create_model("mnasnet_a1",
num_classes=1000,
in_chans=3,
pretrained=False)
if pretrained:
cached_file = '/home/daming/Work/HandDetection/CenterNetUlsB32/src/lib/models/networks/mnasnet_a1.pth'
state_dict = torch.load(cached_file)
load_model(model, state_dict)
return model
def __init__(self, num_classes=1000):
super(FinetuneResnet, self).__init__()
#self.model = models.resnet18(pretrained=False)
#self.fc1 = nn.Linear(512 + 1, 1) # hidden vector + one-hot vector
self.model = geffnet.create_model('mobilenetv3_large_100',
pretrained=False)
self.num_ftrs_1 = self.model.classifier.in_features
self.classifier = nn.Linear(self.num_ftrs_1 + 1,
1) # hidden vector + one-hot vector
self.n_classes = num_classes
def __init__(self,
backbone,
neck=None,
shared_head=None,
rpn_head=None,
bbox_roi_extractor=None,
bbox_head=None,
mask_roi_extractor=None,
mask_head=None,
train_cfg=None,
test_cfg=None,
pretrained=None):
super(TwoStageDetector, self).__init__()
self.eff_backbone_flag = False if 'efficient' not in backbone['type'] else True
if self.eff_backbone_flag == False:
self.backbone = builder.build_backbone(backbone)
else:
self.backbone = geffnet.create_model(backbone['type'],
pretrained=True if pretrained is not None else False,
se=False,
act_layer=backbone['act_cfg']['type'],
norm_layer=norm_cfg[backbone['norm_cfg']['type']][1])
if neck is not None:
self.neck = builder.build_neck(neck)
if shared_head is not None:
self.shared_head = builder.build_shared_head(shared_head)
if rpn_head is not None:
self.rpn_head = builder.build_head(rpn_head)
if bbox_head is not None:
self.bbox_roi_extractor = builder.build_roi_extractor(
bbox_roi_extractor)
self.bbox_head = builder.build_head(bbox_head)
if mask_head is not None:
if mask_roi_extractor is not None:
self.mask_roi_extractor = builder.build_roi_extractor(
mask_roi_extractor)
self.share_roi_extractor = False
else:
self.share_roi_extractor = True
self.mask_roi_extractor = self.bbox_roi_extractor
self.mask_head = builder.build_head(mask_head)
self.train_cfg = train_cfg
self.test_cfg = test_cfg
self.init_weights(pretrained=pretrained)
def __init__(self, train_dataset, val_dataset, logger, hparams):
super().__init__()
self.hparams = hparams
self.train_dataset = train_dataset
self.val_dataset = val_dataset
self.logger = logger
self.model = create_model(self.hparams.model_name,
pretrained=self.hparams.pretrained,
num_classes=self.hparams.num_classes)
self.model.global_pool = GeM()
def get_model():
model = create_model('efficientnet_b0', pretrained=True)
for param in model.parameters():
param.requires_grad = True
fc = nn.Sequential(
OrderedDict([('fc1', nn.Linear(2048, 1000, bias=True)),
('BN1',
nn.BatchNorm2d(1000,
eps=1e-05,
momentum=0.1,
affine=True,
track_running_stats=True)),
('dropout1', nn.Dropout(0.7)),
('fc2', nn.Linear(1000, 512)),
('BN2',
nn.BatchNorm2d(512,
eps=1e-05,
momentum=0.1,
affine=True,
track_running_stats=True)),
('swish1', Swish()), ('dropout2', nn.Dropout(0.5)),
('fc3', nn.Linear(512, 128)),
('BN3',
nn.BatchNorm2d(128,
eps=1e-05,
momentum=0.1,
affine=True,
track_running_stats=True)),
('swish2', Swish()), ('fc4', nn.Linear(128, 3)),
('output', nn.Softmax(dim=1))]))
model.fc = fc
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(),
lr=0.01,
momentum=0.9,
nesterov=True,
weight_decay=0.0001)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)
CHECK_POINT_PATH = '/Users/mac/PycharmProjects/Covid-19_Chest_Xray/weights/EfficientNet_B0_Covid-19.pth'
checkpoint = torch.load(CHECK_POINT_PATH, map_location='cpu')
model.load_state_dict(checkpoint['model_state_dict'])
best_model_wts = copy.deepcopy(model.state_dict())
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
best_loss = checkpoint['best_val_loss']
best_acc = checkpoint['best_val_accuracy']
model.eval()
return model
def __init__(self, backbone, out_dim, n_meta_features=0, load_pretrained=False):
# backbone fait référence au model de base
super(enetv2, self).__init__()# super () vous donne accès aux méthodes d'une superclasse de la sous-classe qui en hérite.
self.n_meta_features = n_meta_features # les meta_features sont initialisés à Zero
self.enet = geffnet.create_model(enet_type.replace('-', '_'), pretrained=load_pretrained)
#create_model est utilisé pour le deploiement d'un modèle, dans les scénarios où vous devez créer un pipeline pour les inférences via plusieurs modèles.
self.dropout = nn.Dropout(0.5)
# Pendant l'entraînement, Dropout() met à zéro au hasard certains des éléments du tenseur d'entrée avec une probabilité 0,5,
# en utilisant des échantillons d'une distribution de Bernoulli
in_ch = self.enet.classifier.in_features #in_feature est le nombre d'input pour notre couche d'entrée
self.myfc = nn.Linear(in_ch, out_dim) # transformation linéaire entre le nombre de notre couche d'entrée et de sortie. ça definit nos couches cachées
self.enet.classifier = nn.Identity() #
def get_model(model_name,
output_size,
pretrained=True,
feature_size=256,
pool=False,
dropout=0.1,
amp=False):
if model_name.startswith('resne'):
m = getattr(models, model_name)
model = m(pretrained=pretrained)
last_num = model.fc.in_features
last = 'fc'
elif model_name.startswith('se'):
model = pretrainedmodels.__dict__[model_name](
num_classes=1000, pretrained='imagenet' if pretrained else None)
model.dropout = None
last_num = model.last_linear.in_features
last = 'last_linear'
elif model_name.startswith('densenet'):
m = getattr(models, model_name)
model = m(pretrained=pretrained)
last_num = model.classifier.in_features
last = 'classifier'
elif model_name.startswith('my_densenet'):
model = models.DenseNet(32,
block_config=(6, 12, 32),
num_init_features=64,
num_classes=output_size)
elif model_name in [
'efficientnet_b0', 'efficientnet_b1', 'efficientnet_b2',
'efficientnet_b3'
] or model_name.startswith('tf_'):
model = geffnet.create_model(model_name, pretrained=pretrained)
last_num = model.classifier.in_features
last = 'classifier'
else:
raise ValueError('no model named ' + model_name)
if not pool:
setattr(model, last, nn.Linear(last_num, feature_size))
else:
setattr(
model, last,
nn.Sequential(NoopAddDim(1), nn.AdaptiveMaxPool1d(feature_size),
NoopSqueezeDim(1)))
if amp:
model = AutocastModule(model)
sq = OrderedDict([('wso', SplitCT()), ('base_model', model),
('drop_out', nn.Dropout(dropout)),
('last_linear', nn.Linear(feature_size, output_size))])
return nn.Sequential(sq)
def make_geffnet(input_size, output_size):
try:
import geffnet
except ImportError:
print('geffnet is not installed!')
print('> pip install geffnet')
raise
# pytorch is CHW
model = geffnet.create_model(name,
num_classes=output_size,
in_chans=input_size[0])
model.__doc__ = delayed_geffnet.__doc__
return model
def __init__(self,
model_name,
pretrained=True,
out_indices=(2, 3, 4, 5, 6),
style='pytorch',
frozen_stages=-1,
norm_eval=True):
super(EfficientNet, self).__init__()
self.out_indices = out_indices
self.style = style
self.frozen_stages = frozen_stages
self.norm_eval = norm_eval
self.model = geffnet.create_model(model_name,pretrained=pretrained)
self._freeze_stages()
def build_encoder(name: str,
weights=None,
freeze=True) -> Tuple[nn.Module, int]:
model = geffnet.create_model(name, pretrained=False)
if weights is not None:
weights = torch.load(weights)
model.load_state_dict(weights)
layers = list(model.children())
model = nn.Sequential(*layers[:-1])
if freeze:
for p in model.parameters():
p.requires_grad_(False)
model.eval()
x = torch.rand(1, 3, 224, 224)
x = model(x)
return model, x.shape[1]
def __init__(self,
enet_type,
out_dim,
drop_nums=1,
pretrained=False,
metric_strategy=False):
super(Effnet, self).__init__()
self.enet = geffnet.create_model(enet_type, pretrained=pretrained)
self.dropouts = nn.ModuleList(
[nn.Dropout(0.5) for _ in range(drop_nums)])
in_ch = self.enet.classifier.in_features
self.fc = nn.Linear(in_ch, 512)
self.swish = Swish_module()
self.metric_classify = ArcMarginProduct_subcenter(512, out_dim)
self.classify = nn.Linear(in_ch, out_dim)
self.enet.classifier = nn.Identity()
self.metric_strategy = metric_strategy
def __init__(self):
super().__init__()
self.model = geffnet.create_model('tf_efficientnet_lite4', pretrained=True)
self.model.global_pool = nn.AdaptiveAvgPool2d((1,1))
in_ch = self.model.classifier.in_features
self.myfc = nn.Sequential(
nn.Dropout(0.17418),
nn.Linear(in_ch, 785),
nn.BatchNorm1d(785),
Swish_module(),
nn.Dropout(0.12),
nn.Linear(785, 1038),
nn.BatchNorm1d(1038),
Swish_module(),
nn.Dropout(0.3763722),
nn.Linear(1038, 24)
)
self.model.classifier = nn.Identity()
def __init__(self, model_name, outputs, exportable=False, **kwargs):
super(EffNet, self).__init__()
self.outputs = outputs
if exportable:
import geffnet
geffnet.config.set_exportable(True)
model = geffnet.create_model(model_name, **kwargs)
else:
model = torch.hub.load("rwightman/gen-efficientnet-pytorch",
model_name, **kwargs)
self.conv_stem = model.conv_stem
self.bn1 = model.bn1
self.act1 = model.act1
for j in range(7):
self.add_module("block{}".format(j + 1),
getattr(model.blocks, "{}".format(j)))
def __init__(self):
super().__init__()
self.model = geffnet.create_model('tf_efficientnet_lite4', pretrained=True)
self.model.global_pool = nn.AdaptiveAvgPool2d((1,1))
in_ch = self.model.classifier.in_features
self.myfc = nn.Sequential(
nn.Dropout(0.29),
nn.Linear(in_ch, 1307),
nn.BatchNorm1d(1307),
nn.ReLU(),
nn.Dropout(0.2),
nn.Linear(1307, 1307),
nn.BatchNorm1d(1307),
nn.ReLU(),
nn.Dropout(0.2),
nn.Linear(1307, 24)
)
self.model.classifier = nn.Identity()
def get_model():
model = create_model('efficientnet_b0', pretrained=True)
model.classifier = nn.Linear(model.classifier.in_features, 120)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(),
lr=0.001,
momentum=0.9,
nesterov=True,
weight_decay=0.0001)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)
CHECK_POINT_PATH = '/PATH/TO/CHECKPOINT/EfficientNet_B0_SGD.pth'
checkpoint = torch.load(CHECK_POINT_PATH, map_location='cpu')
model.load_state_dict(checkpoint['model_state_dict'])
best_model_wts = copy.deepcopy(model.state_dict())
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
best_loss = checkpoint['best_val_loss']
best_acc = checkpoint['best_val_accuracy']
model.eval()
return model
def __init__(self,
backbone,
neck=None,
rpn_head=None,
roi_head=None,
train_cfg=None,
test_cfg=None,
pretrained=None):
super(EfficientRCNN, self).__init__()
self.eff_backbone_flag = False if 'efficient' not in backbone[
'type'] else True
if self.eff_backbone_flag == False:
self.backbone = build_backbone(backbone)
else:
self.backbone = geffnet.create_model(
backbone['type'],
pretrained=True if pretrained is not None else False,
norm_layer=norm_cfg[backbone['norm_cfg']['type']][1])
if neck is not None:
self.neck = build_neck(neck)
if rpn_head is not None:
rpn_train_cfg = train_cfg.rpn if train_cfg is not None else None
rpn_head_ = rpn_head.copy()
rpn_head_.update(train_cfg=rpn_train_cfg, test_cfg=test_cfg.rpn)
self.rpn_head = build_head(rpn_head_)
if roi_head is not None:
# update train and test cfg here for now
# TODO: refactor assigner & sampler
rcnn_train_cfg = train_cfg.rcnn if train_cfg is not None else None
roi_head.update(train_cfg=rcnn_train_cfg)
roi_head.update(test_cfg=test_cfg.rcnn)
self.roi_head = build_head(roi_head)
self.train_cfg = train_cfg
self.test_cfg = test_cfg
self.init_weights(pretrained=pretrained)
def get_model():
#### Config load From Model had use Training
## can change model apply with model use train
# List model is can use with had already train: efficientnet_es, efficientnet_b3
use_model = 'efficientnet_b3'
model = create_model(use_model, pretrained=True)
# Set total fruit index :130 +1
n_class = 131
model.classifier = nn.Linear(model.classifier.in_features, n_class)
criterion = nn.CrossEntropyLoss()
#optimizer = optim.Adam(model.parameters(), lr=0.001, weight_decay=0.0001)
optimizer = optim.SGD(model.parameters(),
lr=0.001,
momentum=0.9,
nesterov=True,
weight_decay=0.0001)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)
CHECK_POINT_PATH = ''
# Config path load check point base on current model using train
if use_model == 'efficientnet_b3':
CHECK_POINT_PATH = 'D:/Learns/check_point/BK_1_CHECKPOINT_EFFICIENTNET_B3_TRAIN_SGD_10_07_20.pth'
if use_model == 'efficientnet_es':
CHECK_POINT_PATH = 'D:/Learns/check_point/CHECKPOINT_MODEL_EFFICIENTNET_ES_TRAIN_SGD.pth'
print('Check point path: ' + CHECK_POINT_PATH)
# load check point
checkpoint = torch.load(CHECK_POINT_PATH, map_location='cpu')
model.load_state_dict(checkpoint['model_state_dict'])
best_model_wts = copy.deepcopy(model.state_dict())
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
best_loss = checkpoint['best_val_loss']
best_acc = checkpoint['best_val_accuracy']
model.eval()
return model
def main():
args = parser.parse_args()
print(args)
if args.img_size is None:
args.img_size, args.crop_pct = get_image_size_crop_pct(args.model)
if not args.checkpoint and not args.pretrained:
args.pretrained = True
if args.torchscript:
geffnet.config.set_scriptable(True)
# create model
model = geffnet.create_model(args.model,
num_classes=args.num_classes,
in_chans=3,
pretrained=args.pretrained,
checkpoint_path=args.checkpoint)
if args.torchscript:
torch.jit.optimized_execution(True)
model = torch.jit.script(model)
print('Model %s created, param count: %d' %
(args.model, sum([m.numel() for m in model.parameters()])))
data_config = resolve_data_config(model, args)
criterion = nn.CrossEntropyLoss()
if not args.no_cuda:
if args.num_gpu > 1:
model = torch.nn.DataParallel(model,
device_ids=list(range(
args.num_gpu))).cuda()
else:
model = model.cuda()
criterion = criterion.cuda()
if args.tune:
model.eval()
model.fuse_model()
conf_yaml = "conf_" + args.model + ".yaml"
from lpot import Quantization
quantizer = Quantization(conf_yaml)
q_model = quantizer(model)
exit(0)
valdir = os.path.join(args.data, 'val')
loader = create_loader(Dataset(valdir, load_bytes=args.tf_preprocessing),
input_size=data_config['input_size'],
batch_size=args.batch_size,
use_prefetcher=not args.no_cuda,
interpolation=data_config['interpolation'],
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
crop_pct=data_config['crop_pct'],
tensorflow_preprocessing=args.tf_preprocessing)
batch_time = AverageMeter()
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
model.eval()
model.fuse_model()
if args.int8:
from lpot.utils.pytorch import load
new_model = load(
os.path.abspath(os.path.expanduser(args.tuned_checkpoint)), model)
else:
new_model = model
with torch.no_grad():
for i, (input, target) in enumerate(loader):
if i >= args.warmup_iterations:
start = time.time()
if not args.no_cuda:
target = target.cuda()
input = input.cuda()
# compute output
output = new_model(input)
loss = criterion(output, target)
# measure accuracy and record loss
prec1, prec5 = accuracy(output.data, target, topk=(1, 5))
losses.update(loss.item(), input.size(0))
top1.update(prec1.item(), input.size(0))
top5.update(prec5.item(), input.size(0))
if i >= args.warmup_iterations:
# measure elapsed time
batch_time.update(time.time() - start)
if i % args.print_freq == 0:
print(
'Test: [{0}/{1}]\t'
'Time {batch_time.val:.3f} ({batch_time.avg:.3f}, {rate_avg:.3f}/s) \t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
i,
len(loader),
batch_time=batch_time,
rate_avg=input.size(0) / batch_time.avg,
loss=losses,
top1=top1,
top5=top5))
if args.iterations > 0 and i >= args.iterations + args.warmup_iterations - 1:
break
print('Batch size = %d' % args.batch_size)
if args.batch_size == 1:
print('Latency: %.3f ms' % (batch_time.avg * 1000))
print('Throughput: %.3f images/sec' %
(args.batch_size / batch_time.avg))
print('Accuracy: {top1:.5f} Accuracy@5 {top5:.5f}'.format(
top1=(top1.avg / 100), top5=(top5.avg / 100)))
def main():
args = parser.parse_args()
args.pretrained = True
if args.checkpoint:
args.pretrained = False
print("==> Creating PyTorch {} model".format(args.model))
# NOTE exportable=True flag disables autofn/jit scripted activations and uses Conv2dSameExport layers
# for models using SAME padding
model = geffnet.create_model(
args.model,
num_classes=args.num_classes,
in_chans=3,
pretrained=args.pretrained,
checkpoint_path=args.checkpoint,
exportable=True)
model.eval()
example_input = torch.randn((args.batch_size, 3, args.img_size or 224, args.img_size or 224), requires_grad=True)
# Run model once before export trace, sets padding for models with Conv2dSameExport. This means
# that the padding for models with Conv2dSameExport (most models with tf_ prefix) is fixed for
# the input img_size specified in this script.
# Opset >= 11 should allow for dynamic padding, however I cannot get it to work due to
# issues in the tracing of the dynamic padding or errors attempting to export the model after jit
# scripting it (an approach that should work). Perhaps in a future PyTorch or ONNX versions...
model(example_input)
print("==> Exporting model to ONNX format at '{}'".format(args.output))
input_names = ["input0"]
output_names = ["output0"]
dynamic_axes = {'input0': {0: 'batch'}, 'output0': {0: 'batch'}}
if args.dynamic_size:
dynamic_axes['input0'][2] = 'height'
dynamic_axes['input0'][3] = 'width'
if args.aten_fallback:
export_type = torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK
else:
export_type = torch.onnx.OperatorExportTypes.ONNX
torch_out = torch.onnx._export(
model, example_input, args.output, export_params=True, verbose=True, input_names=input_names,
output_names=output_names, keep_initializers_as_inputs=args.keep_init, dynamic_axes=dynamic_axes,
opset_version=args.opset, operator_export_type=export_type)
print("==> Loading and checking exported model from '{}'".format(args.output))
onnx_model = onnx.load(args.output)
onnx.checker.check_model(onnx_model) # assuming throw on error
print("==> Passed")
if args.keep_init and args.aten_fallback:
import caffe2.python.onnx.backend as onnx_caffe2
# Caffe2 loading only works properly in newer PyTorch/ONNX combos when
# keep_initializers_as_inputs and aten_fallback are set to True.
print("==> Loading model into Caffe2 backend and comparing forward pass.".format(args.output))
caffe2_backend = onnx_caffe2.prepare(onnx_model)
B = {onnx_model.graph.input[0].name: x.data.numpy()}
c2_out = caffe2_backend.run(B)[0]
np.testing.assert_almost_equal(torch_out.data.numpy(), c2_out, decimal=5)
print("==> Passed")