Spaces:
Running
Running
File size: 9,695 Bytes
63f3cf2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 |
# -*- coding: UTF-8 -*-
'''=================================================
@Project -> File pram -> train
@IDE PyCharm
@Author fx221@cam.ac.uk
@Date 29/01/2024 14:26
=================================================='''
import argparse
import os
import os.path as osp
import torch
import torchvision.transforms.transforms as tvt
import yaml
import torch.utils.data as Data
import torch.multiprocessing as mp
import torch.distributed as dist
from nets.segnet import SegNet
from nets.segnetvit import SegNetViT
from dataset.utils import collect_batch
from dataset.get_dataset import compose_datasets
from tools.common import torch_set_gpu
from trainer import Trainer
from nets.sfd2 import ResNet4x, DescriptorCompressor
from nets.superpoint import SuperPoint
torch.set_grad_enabled(True)
parser = argparse.ArgumentParser(description='PRAM', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--config', type=str, required=True, help='config of specifications')
parser.add_argument('--landmark_path', type=str, default=None, help='path of landmarks')
def load_feat_network(config):
if config['feature'] == 'spp':
net = SuperPoint(config={
'weight_path': '/scratches/flyer_2/fx221/Research/Code/third_weights/superpoint_v1.pth',
}).eval()
elif config['feature'] == 'resnet4x':
net = ResNet4x(inputdim=3, outdim=128)
net.load_state_dict(
torch.load('weights/sfd2_20230511_210205_resnet4x.79.pth', map_location='cpu')['state_dict'],
strict=True)
net.eval()
else:
print('Please input correct feature {:s}'.format(config['feature']))
net = None
if config['feat_dim'] != 128:
desc_compressor = DescriptorCompressor(inputdim=128, outdim=config['feat_dim']).eval()
if config['feat_dim'] == 64:
desc_compressor.load_state_dict(
torch.load('weights/20230511_210205_resnet4x_B6_R512_I3_O128_pho_resnet4x_e79_to_O64.pth',
map_location='cpu'),
strict=True)
elif config['feat_dim'] == 32:
desc_compressor.load_state_dict(
torch.load('weights/20230511_210205_resnet4x_B6_R512_I3_O128_pho_resnet4x_e79_to_O32.pth',
map_location='cpu'),
strict=True)
else:
desc_compressor = None
else:
desc_compressor = None
return net, desc_compressor
def get_model(config):
desc_dim = 256 if config['feature'] == 'spp' else 128
if config['use_mid_feature']:
desc_dim = 256
model_config = {
'network': {
'descriptor_dim': desc_dim,
'n_layers': config['layers'],
'ac_fn': config['ac_fn'],
'norm_fn': config['norm_fn'],
'n_class': config['n_class'],
'output_dim': config['output_dim'],
'with_cls': config['with_cls'],
'with_sc': config['with_sc'],
'with_score': config['with_score'],
}
}
if config['network'] == 'segnet':
model = SegNet(model_config.get('network', {}))
config['with_cls'] = False
elif config['network'] == 'segnetvit':
model = SegNetViT(model_config.get('network', {}))
config['with_cls'] = False
else:
raise 'ERROR! {:s} model does not exist'.format(config['network'])
if config['local_rank'] == 0:
if config['weight_path'] is not None:
state_dict = torch.load(osp.join(config['save_path'], config['weight_path']), map_location='cpu')['model']
model.load_state_dict(state_dict, strict=True)
print('Load weight from {:s}'.format(osp.join(config['save_path'], config['weight_path'])))
if config['resume_path'] is not None and not config['eval']: # only for training
model.load_state_dict(
torch.load(osp.join(config['save_path'], config['resume_path']), map_location='cpu')['model'],
strict=True)
print('Load resume weight from {:s}'.format(osp.join(config['save_path'], config['resume_path'])))
return model
def setup(rank, world_size):
os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '12355'
# initialize the process group
dist.init_process_group("nccl", rank=rank, world_size=world_size)
def train_DDP(rank, world_size, model, config, train_set, test_set, feat_model, img_transforms):
print('In train_DDP..., rank: ', rank)
torch.cuda.set_device(rank)
device = torch.device(f'cuda:{rank}')
if feat_model is not None:
feat_model.to(device)
model.to(device)
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
setup(rank=rank, world_size=world_size)
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[rank])
train_sampler = torch.utils.data.distributed.DistributedSampler(train_set,
shuffle=True,
rank=rank,
num_replicas=world_size,
drop_last=True, # important?
)
train_loader = torch.utils.data.DataLoader(train_set,
batch_size=config['batch_size'] // world_size,
num_workers=config['workers'] // world_size,
# num_workers=1,
pin_memory=True,
# persistent_workers=True,
shuffle=False, # must be False
drop_last=True,
collate_fn=collect_batch,
prefetch_factor=4,
sampler=train_sampler)
config['local_rank'] = rank
if rank == 0:
test_set = test_set
else:
test_set = None
trainer = Trainer(model=model, train_loader=train_loader, feat_model=feat_model, eval_loader=test_set,
config=config, img_transforms=img_transforms)
trainer.train()
if __name__ == '__main__':
args = parser.parse_args()
with open(args.config, 'rt') as f:
config = yaml.load(f, Loader=yaml.Loader)
torch_set_gpu(gpus=config['gpu'])
if config['local_rank'] == 0:
print(config)
if config['feature'] == 'spp':
img_transforms = None
else:
img_transforms = []
img_transforms.append(tvt.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]))
img_transforms = tvt.Compose(img_transforms)
feat_model, desc_compressor = load_feat_network(config=config)
dataset = config['dataset']
if config['eval'] or config['loc']:
if not config['online']:
from localization.loc_by_rec_eval import loc_by_rec_eval
test_set = compose_datasets(datasets=dataset, config=config, train=False, sample_ratio=1)
config['n_class'] = test_set.n_class
model = get_model(config=config)
loc_by_rec_eval(rec_model=model.cuda().eval(),
loader=test_set,
local_feat=feat_model.cuda().eval(),
config=config, img_transforms=img_transforms)
else:
from localization.loc_by_rec_online import loc_by_rec_online
model = get_model(config=config)
loc_by_rec_online(rec_model=model.cuda().eval(),
local_feat=feat_model.cuda().eval(),
config=config, img_transforms=img_transforms)
exit(0)
train_set = compose_datasets(datasets=dataset, config=config, train=True, sample_ratio=None)
if config['do_eval']:
test_set = compose_datasets(datasets=dataset, config=config, train=False, sample_ratio=None)
else:
test_set = None
config['n_class'] = train_set.n_class
model = get_model(config=config)
if not config['with_dist'] or len(config['gpu']) == 1:
config['with_dist'] = False
model = model.cuda()
train_loader = Data.DataLoader(dataset=train_set,
shuffle=True,
batch_size=config['batch_size'],
drop_last=True,
collate_fn=collect_batch,
num_workers=config['workers'])
if test_set is not None:
test_loader = Data.DataLoader(dataset=test_set,
shuffle=False,
batch_size=1,
drop_last=False,
collate_fn=collect_batch,
num_workers=4)
else:
test_loader = None
trainer = Trainer(model=model, train_loader=train_loader, feat_model=feat_model, eval_loader=test_loader,
config=config, img_transforms=img_transforms)
trainer.train()
else:
mp.spawn(train_DDP, nprocs=len(config['gpu']),
args=(len(config['gpu']), model, config, train_set, test_set, feat_model, img_transforms),
join=True)
|