def make_vgg(cfg_id,
batch_norm,
pretrained,
url_id,
incl_fcs=True,
num_classes=1000):
model = VGG(cfgs[cfg_id],
batch_norm=batch_norm,
incl_fcs=incl_fcs,
num_classes=num_classes)
if pretrained:
pretrained_model = model_zoo.load_url(model_urls[url_id])
for k, v in fc_mapping.items():
pretrained_model[k] = pretrained_model.pop(v)
if num_classes != 1000:
del pretrained_model['classifier.weight']
del pretrained_model['classifier.bias']
try:
load_state_dict(model, pretrained_model)
except KeyError:
from jacinle.logging import get_logger
logger = get_logger(__file__)
logger.exception('test')
pass # Intentionally ignore the key error.
return model
def make_resnet(net_id, pretrained, incl_gap=True, num_classes=1000):
model = ResNet(*cfgs[net_id], incl_gap=incl_gap, num_classes=num_classes)
if pretrained:
pretrained_model = model_zoo.load_url(model_urls[net_id])
if num_classes != 1000:
del pretrained_model['fc.weight']
del pretrained_model['fc.bias']
try:
load_state_dict(model, pretrained_model)
except KeyError:
pass # Intentionally ignore the key error.
return model
def load_checkpoint(self, filename):
if osp.isfile(filename):
model = self._model
if isinstance(model, nn.DataParallel):
model = model.module
try:
checkpoint = torch.load(filename)
load_state_dict(model, checkpoint['model'])
self._optimizer.load_state_dict(checkpoint['optimizer'])
logger.critical('Checkpoint loaded: {}.'.format(filename))
return checkpoint['extra']
except Exception:
logger.exception('Error occurred when load checkpoint "{}".'.format(filename))
else:
logger.warning('No checkpoint found at specified position: "{}".'.format(filename))
return None