def main(args):
visenv_name = args.dataset
exp_dir = os.path.join(args.save_path, args.dataset)
model_dir, log_dir = get_model_log_path(exp_dir, visenv_name)
stdout_file = os.path.join(log_dir, f'stdout_{time_str()}.txt')
save_model_path = os.path.join(model_dir, 'ckpt_max.pth')
os.environ['CUDA_VISIBLE_DEVICES'] = args.device
if args.redirector:
print('redirector stdout')
ReDirectSTD(stdout_file, 'stdout', False)
pprint.pprint(OrderedDict(args.__dict__))
print('-' * 60)
print(f'use GPU{args.device} for training')
print(f'train set: {args.dataset} {args.train_split}, test set: {args.valid_split}')
#train_tsfm, valid_tsfm = get_transform(args)
train_tsfm, valid_tsfm = get_transform(args)
print(train_tsfm)
#train_set = AttrDataset(args=args, split=args.train_split, transform=train_tsfm)
train_set = AttrDataset(args=args, split=args.train_split, transform=train_tsfm, target_transform=None, Type='train' )
train_loader = DataLoader(
dataset=train_set,
batch_size=args.batchsize,
shuffle=True,
num_workers=4,
drop_last=True,
pin_memory=True,
)
#valid_set = AttrDataset(args=args, split=args.valid_split, transform=valid_tsfm)
valid_set = AttrDataset(args=args, split=args.valid_split, transform=valid_tsfm, target_transform=None, Type='val')
valid_loader = DataLoader(
dataset=valid_set,
batch_size=args.batchsize,
shuffle=False,
num_workers=4,
drop_last=True,
pin_memory=True,
)
print(f'{args.train_split} set: {len(train_loader.dataset)}, '
f'{args.valid_split} set: {len(valid_loader.dataset)}, '
f'attr_num : {train_set.attr_num}')
labels = train_set.label
sample_weight = labels.mean(0)
#backbone = resnet50()
if args.model_name == 'resnet50':
backbone = resnet50()
if args.model_name == 'resnet18':
backbone = resnet18()
if args.model_name == 'resnet18_stn':
backbone = resnet18_stn()
if args.model_name == 'resnet_depth':
backbone = resnet_depth()
if args.model_name == 'resnet18_transformer':
backbone = resnet18_transformer()
if args.model_name == 'resnet50_dynamic_se':
backbone = resnet50_dynamic_se()
if args.model_name == 'resnet18_dynamic_se':
backbone = resnet18_dynamic_se()
if args.model_name == 'resnet18_replace_se':
backbone = resnet18_replace_se()
if args.model_name == 'resnet18_se':
backbone = resnet18_se()
if args.model_name == 'resnet34':
backbone = resnet34()
if args.model_name == 'resnet18_group_se':
backbone = resnet18_group_se()
if args.model_name == 'resnet18_vit':
backbone = resnet18_vit()
if args.model_name == 'resnet18_vit_v2':
backbone = resnet18_vit_v2()
if args.model_name == 'resnet18_vit_v3':
backbone = resnet18_vit_v3()
if args.model_name == 'resnet18_vit_v5':
backbone = resnet18_vit_v5()
if args.model_name == 'resnet18_energy_vit':
backbone = resnet18_energy_vit()
if args.model_name == 'resnet18_vit_split':
backbone = resnet18_vit_split(num_classes = train_set.attr_num)
if args.model_name == 'inception_self':
backbone = inception_self()
if args.model_name == 'spatial_modulator':
backbone = spatial_modulator()
if args.model_name == 'fusion_concat':
backbone = fusion_concat()
print('have generated the model')
classifier = BaseClassifier(nattr=train_set.attr_num)
model = FeatClassifier(backbone, classifier)
print('Number of model parameters: {}'.format(
sum([p.data.nelement() for p in model.parameters()])))
print('')
#if torch.cuda.is_available():
model = torch.nn.DataParallel(model).cuda()
#for k, v in model.state_dict().items():
# print(k)
'''
model_dict = {}
state_dict = model.state_dict()
pretrain_dict = torch.load('/home/pengqy/paper/resnet18_2/PETA/PETA/img_model/ckpt_max.pth')['state_dicts']
for k, v in pretrain_dict.items():
# print('%%%%% ', k)
if k in state_dict:
if k.startswith('module.backbone.conv1'):
#pdb.set_trace()
model_dict[k] = v
elif k.startswith('module.backbone.bn1'):
model_dict[k] = v
elif k.startswith('module.backbone.layer'):
model_dict[k] = v
elif k.startswith('module.classifier'):
model_dict[k] = v
#elif k.startswith('module.backbone.spa_conv_0'):
# model_dict[k] = v
#elif k.startswith('module.backbone.spa_bn_0'):
# model_dict[k] = v
#elif k.startswith('module.classifier'):
# model_dict[k] = v
#elif k.startswith('module.classifier'):
# model_dict[k] = v
#pdb.set_trace()
state_dict.update(model_dict)
model.load_state_dict(state_dict)
'''
criterion = CEL_Sigmoid(sample_weight)
param_groups = [{'params': model.module.finetune_params(), 'lr':0.01},
{'params': model.module.fresh_params(), 'lr':0.1}]
optimizer = torch.optim.SGD(param_groups, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=False)
lr_scheduler = ReduceLROnPlateau(optimizer, factor=0.1, patience=4)
loss = args.loss
best_metric, epoch = trainer(epoch=args.train_epoch,
model=model,
train_loader=train_loader,
valid_loader=valid_loader,
criterion=criterion,
optimizer=optimizer,
lr_scheduler=lr_scheduler,
path=save_model_path,
loss =loss)
print(f'{visenv_name}, best_metrc : {best_metric} in epoch{epoch}')
def __init__(self):
device = torch.device('cpu')
FORCE_TO_CPU = True
parser = argument_parser()
args = parser.parse_args(['PETA', '--model=dpn107'])
visenv_name = 'PETA'
exp_dir = os.path.join('exp_result', visenv_name)
model_dir, log_dir = get_model_log_path(exp_dir, visenv_name)
stdout_file = os.path.join(log_dir, f'stdout_{time_str()}.txt')
save_model_path = os.path.join(
model_dir, 'ckpt_max_e0384293_2020-09-17_18-35-21.pth')
if args.redirector:
print('redirector stdout')
ReDirectSTD(stdout_file, 'stdout', False)
pprint.pprint(OrderedDict(args.__dict__))
print('-' * 60)
print(f'use GPU{args.device} for training')
_, predict_tsfm = get_transform(args)
valid_set = AttrDataset(args=args,
split=args.valid_split,
transform=predict_tsfm)
args.att_list = valid_set.attr_id
backbone = getattr(sys.modules[__name__], args.model)()
if "dpn68" in args.model:
net_parameter = 832
elif "dpn" in args.model:
net_parameter = 2688
elif "densenet" in args.model:
net_parameter = 1024
else:
net_parameter = 2048
classifier = BaseClassifier(netpara=net_parameter,
nattr=valid_set.attr_num)
model = FeatClassifier(backbone, classifier)
if torch.cuda.is_available() and not FORCE_TO_CPU:
model = torch.nn.DataParallel(model).cuda()
ckpt = torch.load(save_model_path)
print(f'Model is served with GPU ')
else:
model = torch.nn.DataParallel(model)
ckpt = torch.load(save_model_path,
map_location=torch.device('cpu'))
print(f'Model is served with CPU ')
model.load_state_dict(ckpt['state_dicts'])
model.eval()
# from torchsummary import summary
# summary(model, input_size=(3, 256, 192))
print('Total number of parameters: ',
sum(p.numel() for p in model.parameters() if p.requires_grad))
self.args = args
self.predict_tsfm = predict_tsfm
self.model = model
def main(args):
visenv_name = args.dataset
exp_dir = os.path.join(args.save_path, args.dataset)
model_dir, log_dir = get_model_log_path(exp_dir, visenv_name)
stdout_file = os.path.join(log_dir, f'stdout_{time_str()}.txt')
save_model_path = os.path.join(model_dir, 'ckpt_max.pth')
os.environ['CUDA_VISIBLE_DEVICES'] = args.device
if args.redirector:
print('redirector stdout')
ReDirectSTD(stdout_file, 'stdout', False)
pprint.pprint(OrderedDict(args.__dict__))
print('-' * 60)
print(f'use GPU{args.device} for training')
print(
f'train set: {args.dataset} {args.train_split}, test set: {args.valid_split}'
)
#train_tsfm, valid_tsfm = get_transform(args)
train_tsfm, valid_tsfm = get_transform(args)
print(train_tsfm)
#train_set = AttrDataset(args=args, split=args.train_split, transform=train_tsfm)
train_set = AttrDataset(args=args,
split=args.train_split,
transform=train_tsfm)
train_loader = DataLoader(
dataset=train_set,
batch_size=args.batchsize,
shuffle=True,
num_workers=4,
pin_memory=True,
)
#valid_set = AttrDataset(args=args, split=args.valid_split, transform=valid_tsfm)
valid_set = AttrDataset(args=args,
split=args.valid_split,
transform=valid_tsfm)
valid_loader = DataLoader(
dataset=valid_set,
batch_size=args.batchsize,
shuffle=False,
num_workers=4,
pin_memory=True,
)
print(f'{args.train_split} set: {len(train_loader.dataset)}, '
f'{args.valid_split} set: {len(valid_loader.dataset)}, '
f'attr_num : {train_set.attr_num}')
labels = train_set.label
sample_weight = labels.mean(0)
#backbone = resnet50()
if args.model_name == 'resnet50':
backbone = resnet50()
if args.model_name == 'resnet18':
backbone = resnet18()
if args.model_name == 'resnet50_dynamic_se':
backbone = resnet50_dynamic_se()
if args.model_name == 'resnet18_dynamic_se':
backbone = resnet18_dynamic_se()
if args.model_name == 'resnet18_replace_se':
backbone = resnet18_replace_se()
if args.model_name == 'resnet18_se':
backbone = resnet18_se()
if args.model_name == 'resnet34':
backbone = resnet34()
if args.model_name == 'acnet':
backbone = resnet18_acnet(num_classes=train_set.attr_num)
print('have generated the model')
classifier = BaseClassifier(nattr=train_set.attr_num)
classifier_depth = BaseClassifier(nattr=train_set.attr_num)
model = FeatClassifier(backbone, classifier)
print('Number of model parameters: {}'.format(
sum([p.data.nelement() for p in model.parameters()])))
print('')
#if torch.cuda.is_available():
model = torch.nn.DataParallel(model).cuda()
criterion = CEL_Sigmoid(sample_weight)
param_groups = [{
'params': model.module.finetune_params(),
'lr': args.lr_ft
}, {
'params': model.module.fresh_params(),
'lr': args.lr_new
}]
optimizer = torch.optim.SGD(param_groups,
momentum=args.momentum,
weight_decay=args.weight_decay,
nesterov=False)
lr_scheduler = ReduceLROnPlateau(optimizer, factor=0.1, patience=4)
loss = args.loss
best_metric, epoch = trainer(epoch=args.train_epoch,
model=model,
train_loader=train_loader,
valid_loader=valid_loader,
criterion=criterion,
optimizer=optimizer,
lr_scheduler=lr_scheduler,
path=save_model_path,
loss=loss)
print(f'{visenv_name}, best_metrc : {best_metric} in epoch{epoch}')
