def sib_init(self, params):
model_name = params.model.lower() # wrn
self.num_classes = 64
if params.dataset == "cross" or params.dataset == "miniImagenet":
model_dir = os.path.join(configs.sib_dir, "miniImagenet",
model_name + "_best.pth")
clf_dir = os.path.join(configs.sib_dir, "miniImagenet",
model_name + "_best_clf.pth")
model = SIBWRN(num_classes=self.num_classes)
model = model.cuda()
# Encoder Init
model_dict = model.encoder.state_dict()
checkpoint = torch.load(model_dir)
model_dict.update(checkpoint)
model.encoder.load_state_dict(model_dict)
# Classifier Init
clf_dict = model.classifier.state_dict()
clf_dict.update(torch.load(clf_dir))
model.classifier.load_state_dict(clf_dict)
model.eval()
self.model = model
self.image_size = 80
self.batch_size = 64
self.feat_dim = 640
def sib_init(params, split):
model_name = params.model.lower()
if params.dataset == "cross" or params.dataset == "miniImagenet":
model_dir = os.path.join(configs.sib_dir, "miniImagenet",
model_name + "_best.pth")
model = SIBWRN(num_classes=64)
model = model.cuda()
model_dict = model.encoder.state_dict()
checkpoint = torch.load(model_dir)
model_dict.update(checkpoint)
model.encoder.load_state_dict(model_dict)
model.eval()
if params.dataset == "cross":
loadfile = configs.data_dir['CUB'] + split + '.json'
elif params.dataset == "miniImagenet":
loadfile = configs.data_dir['miniImagenet'] + split + '.json'
image_size = 80
datamgr = SimpleDataManager(image_size, batch_size=64)
data_loader = datamgr.get_data_loader(loadfile, aug=False, num_workers=12)
outfile = '%s/features/%s/%s/%s.hdf5' % (configs.sib_dir, params.dataset,
params.model, split)
dirname = os.path.dirname(outfile)
if not os.path.isdir(dirname):
os.makedirs(dirname)
return model, data_loader, outfile, params
def simpleshot_init(self, params):
if params.dataset == "tiered":
self.num_classes = 351
else:
self.num_classes = 64
model_name = params.model.lower()
if model_name == "resnet10":
model_abbr = "resnet"
elif model_name == "wideres":
model_abbr = "wrn"
if params.dataset == "cross" or params.dataset == "miniImagenet":
# model_dir = "/model/1154027137/ifsl_mini_pretrain/ifsl_mini/ss_" + model_abbr + "_mini.tar"
model_dir = os.path.join(configs.simple_shot_dir, "miniImagenet",
model_name, "model_best.pth.tar")
elif params.dataset == "tiered":
# model_dir = "/model/1154027137/ifsl_tiered_pretrain/ifsl_tiered/ss_" + model_abbr + "_tiered.tar"
model_dir = os.path.join(configs.simple_shot_dir, "tiered",
model_name, "model_best.pth.tar")
def remove_module_from_param_name(params_name_str):
split_str = params_name_str.split(".")[1:]
params_name_str = ".".join(split_str)
return params_name_str
if model_name == "resnet10":
model = resnet10(num_classes=self.num_classes, remove_linear=False)
elif model_name == "wideres":
model = wideres(num_classes=self.num_classes, remove_linear=False)
# model = simple_shot_models.__dict__[model_name](num_classes=self.num_classes, remove_linear=False)
model = model.cuda()
checkpoint = torch.load(model_dir)
model_dict = model.state_dict()
model_params = checkpoint['state_dict']
model_params = {
remove_module_from_param_name(k): v
for k, v in model_params.items()
}
model_params = {
k: v
for k, v in model_params.items() if k in model_dict
}
model_dict.update(model_params)
model.load_state_dict(model_dict)
model.eval()
self.model = model
self.image_size = 84
if self.model_name == "wideres":
self.batch_size = 128
self.feat_dim = 640
else:
self.batch_size = 128
self.feat_dim = 512
def _model_load(model, pretrained_dict):
model_dict = model.state_dict()
# model_dict_keys = [v.replace('module.', '') for v in model_dict.keys() if v.startswith('module.')]
if list(model_dict.keys())[0].startswith('module.'):
pretrained_dict = {'module.' + k: v for k, v in pretrained_dict.items()}
if list(model_dict.keys())[0].startswith('feature.'):
pretrained_dict = {k.replace("features.","feature."): v for k, v in pretrained_dict.items()}
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict.keys()}
print("Weights cannot be loaded:")
print([k for k in model_dict.keys() if k not in pretrained_dict.keys()])
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
def simple_shot_init(params, split):
model_name = params.model.lower()
if params.dataset == "cross" or params.dataset == "miniImagenet":
model_dir = os.path.join(configs.simple_shot_dir, "miniImagenet",
model_name, "model_best.pth.tar")
elif params.dataset == "tiered":
model_dir = os.path.join(configs.simple_shot_dir, "tiered", model_name,
"model_best.pth.tar")
num_classes = 64
if params.dataset == "tiered":
num_classes = 351
model = simple_shot_models.__dict__[model_name](num_classes=num_classes,
remove_linear=True)
model = model.cuda()
# model = torch.nn.DataParallel(model).cuda()
checkpoint = torch.load(model_dir)
model_dict = model.state_dict()
model_params = checkpoint['state_dict']
model_params = {
remove_module_from_param_name(k): v
for k, v in model_params.items()
}
model_params = {k: v for k, v in model_params.items() if k in model_dict}
model_dict.update(model_params)
model.load_state_dict(model_dict)
tiered_mini = False
if params.dataset == "cross":
loadfile = configs.data_dir['CUB'] + split + '.json'
elif params.dataset == "miniImagenet":
loadfile = configs.data_dir['miniImagenet'] + split + '.json'
elif params.dataset == "tiered":
loadfile = split
tiered_mini = True
image_size = 84
datamgr = SimpleDataManager(image_size, batch_size=64)
data_loader = datamgr.get_data_loader(loadfile,
aug=False,
num_workers=12,
tiered_mini=tiered_mini)
outfile = '%s/features/%s/%s/%s.hdf5' % (
configs.simple_shot_dir, params.dataset, params.model, split)
dirname = os.path.dirname(outfile)
if not os.path.isdir(dirname):
os.makedirs(dirname)
return model, data_loader, outfile, params
def feat_init(self, params):
model_name = params.model.lower() # wrn
self.num_classes = 64
if params.dataset == "cross" or params.dataset == "miniImagenet":
model_dir = os.path.join(configs.feat_dir, "miniImagenet",
model_name + "_pre.pth")
# model_dir = "/model/1154027137/ifsl_mini_pretrain/ifsl_mini/ft_" + "wrn" + "_mini.tar"
model = FEATWRN(num_classes=self.num_classes)
model = model.cuda()
model_dict = model.state_dict()
checkpoint = torch.load(model_dir)
model_dict.update(checkpoint['params'])
model.load_state_dict(model_dict)
model.eval()
self.model = model
self.image_size = 84
self.batch_size = 64
self.feat_dim = 640
def cosine_init(params, split):
model_name = params.model.lower()
if params.dataset == "cross" or params.dataset == "miniImagenet":
num_classes = 64
model_dir = os.path.join(configs.cosine_dir, "miniImagenet",
model_name, "max_acc.pth")
elif params.dataset == "tiered":
num_classes = 351
model_dir = os.path.join(configs.cosine_dir, "tiered", model_name,
"max_acc.pth")
if model_name == "resnet10":
feat_dim = 512
elif model_name == "wrn":
feat_dim = 640
model = CosinePretrain(model_name, num_classes, feat_dim)
model = model.cuda()
model_dict = model.encoder.state_dict()
ckpt = torch.load(model_dir)["params"]
model_dict.update(ckpt)
model.encoder.load_state_dict(model_dict)
model.eval()
if params.dataset == "cross":
loadfile = configs.data_dir['CUB'] + split + '.json'
elif params.dataset == "miniImagenet":
loadfile = configs.data_dir['miniImagenet'] + split + '.json'
image_size = 84
datamgr = SimpleDataManager(image_size, batch_size=64)
data_loader = datamgr.get_data_loader(loadfile, aug=False, num_workers=12)
outfile = '%s/features/%s/%s/%s.hdf5' % (
configs.cosine_dir, params.dataset, params.model, split)
dirname = os.path.dirname(outfile)
if not os.path.isdir(dirname):
os.makedirs(dirname)
return model, data_loader, outfile, params