def get_teacher(p):
# Get backbone
if p['backbone'] == 'resnet18':
if p['train_db_name'] in [
'cifar-10', 'cifar-10-d', 'cifar-10-f', 'cifar-20',
'cifar-20-d', 'cifar-20-f'
]:
from models.resnet_cifar import resnet18
backbone = resnet18()
elif p['train_db_name'] in ['stl-10', 'stl-10-d', 'stl-10-f']:
from models.resnet_stl import resnet18
backbone = resnet18()
elif p['train_db_name'] == 'pascal-pretrained':
from models.resnet_pascal import resnet18
backbone = resnet18()
else:
raise NotImplementedError
elif p['backbone'] == 'resnet50':
if 'imagenet' in p['train_db_name']:
from models.resnet_wider import resnet50x1
backbone = resnet50x1()
else:
raise ValueError('Invalid backbone {}'.format(p['backbone']))
from models.models import ClusteringModel
if p['teacher'] == 'selflabel':
assert (p['num_heads'] == 1)
teacher = ClusteringModel(backbone, p['num_classes'], p['num_heads'])
return teacher
def calc_clusters():
json_events = request.get_json()
events, events_type = parse_events(json_events)
config = clustering_configs[events_type]
model = ClusteringModel(**config)
events_labels = model.predict(events)
labled_events = []
for event, label in events_labels:
obj = {'lat': event.lat,
'lon': event.lon,
'ts': event.time,
'event': events_type,
'cid': label}
labled_events.append(obj)
return jsonify(labled_events)
def get_model(p, pretrain_path=None):
# Get backbone
if p['backbone'] == 'resnet18':
if p['train_db_name'] in ['cifar-10', 'cifar-20']:
from models.resnet_cifar import resnet18
backbone = resnet18()
elif p['train_db_name'] == 'stl-10':
from models.resnet_stl import resnet18
backbone = resnet18()
elif p['train_db_name'] in ['sewer']:
from models.resnet_cifar import resnet18
backbone = resnet18()
else:
raise NotImplementedError
elif p['backbone'] == 'resnet50':
if 'imagenet' in p['train_db_name']:
from models.resnet import resnet50
backbone = resnet50()
else:
raise NotImplementedError
else:
raise ValueError('Invalid backbone {}'.format(p['backbone']))
# Setup
if p['setup'] in ['simclr', 'moco']:
from models.models import ContrastiveModel
model = ContrastiveModel(backbone, **p['model_kwargs'])
elif p['setup'] in ['scan', 'selflabel']:
from models.models import ClusteringModel
if p['setup'] == 'selflabel':
assert (p['num_heads'] == 1)
model = ClusteringModel(backbone, p['num_classes'], p['num_heads'])
else:
raise ValueError('Invalid setup {}'.format(p['setup']))
# Load pretrained weights
if pretrain_path is not None and os.path.exists(pretrain_path):
state = torch.load(pretrain_path, map_location='cpu')
if p['setup'] == 'scan': # Weights are supposed to be transfered from contrastive training
missing = model.load_state_dict(state, strict=False)
assert (set(missing[1]) == {
'contrastive_head.0.weight', 'contrastive_head.0.bias',
'contrastive_head.2.weight', 'contrastive_head.2.bias'
} or set(missing[1])
== {'contrastive_head.weight', 'contrastive_head.bias'})
elif p['setup'] == 'selflabel': # Weights are supposed to be transfered from scan
# We only continue with the best head (pop all heads first, then copy back the best head)
model_state = state['model']
all_heads = [k for k in model_state.keys() if 'cluster_head' in k]
best_head_weight = model_state['cluster_head.%d.weight' %
(state['head'])]
best_head_bias = model_state['cluster_head.%d.bias' %
(state['head'])]
for k in all_heads:
model_state.pop(k)
model_state['cluster_head.0.weight'] = best_head_weight
model_state['cluster_head.0.bias'] = best_head_bias
missing = model.load_state_dict(model_state, strict=True)
else:
raise NotImplementedError
elif pretrain_path is not None and not os.path.exists(pretrain_path):
raise ValueError(
'Path with pre-trained weights does not exist {}'.format(
pretrain_path))
else:
pass
return model
