def main():
start_epoch = 0
args = parse_args()
# Use GPU
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu if args.gpu != '' else str(
opt.gpu_id)
use_gpu = torch.cuda.is_available() and (args.gpu != ''
or int(opt.gpu_id)) >= 0
if not os.path.isdir(opt.checkpoint):
os.makedirs(opt.checkpoint)
# Data
print('==> Preparing dataset')
input_dim = opt.input_size
train_transformer = TrainTransform(size=input_dim)
test_transformer = TestTransform(size=input_dim)
try:
if isinstance(opt.trainset, list):
datalist = []
for dataset, freq, max_skip in zip(opt.trainset, opt.datafreq,
opt.max_skip):
ds = DATA_CONTAINER[dataset](
train=True,
sampled_frames=opt.sampled_frames,
transform=train_transformer,
max_skip=max_skip,
samples_per_video=opt.samples_per_video)
datalist += [ds] * freq
trainset = data.ConcatDataset(datalist)
else:
max_skip = opt.max_skip[0] if isinstance(opt.max_skip,
list) else opt.max_skip
trainset = DATA_CONTAINER[opt.trainset](
train=True,
sampled_frames=opt.sampled_frames,
transform=train_transformer,
max_skip=max_skip,
samples_per_video=opt.samples_per_video)
except KeyError as ke:
print(
'[ERROR] invalide dataset name is encountered. The current acceptable datasets are:'
)
print(list(DATA_CONTAINER.keys()))
exit()
testset = DATA_CONTAINER[opt.valset](train=False,
transform=test_transformer,
samples_per_video=1)
trainloader = data.DataLoader(trainset,
batch_size=opt.train_batch,
shuffle=True,
num_workers=opt.workers,
collate_fn=multibatch_collate_fn)
testloader = data.DataLoader(testset,
batch_size=1,
shuffle=False,
num_workers=opt.workers,
collate_fn=multibatch_collate_fn)
# Model
print("==> creating model")
net = STAN(opt.keydim, opt.valdim)
print(' Total params: %.2fM' %
(sum(p.numel() for p in net.parameters()) / 1000000.0))
net.eval()
if use_gpu:
net = net.cuda()
# set training parameters
for p in net.parameters():
p.requires_grad = True
criterion = None
celoss = cross_entropy_loss
if opt.loss == 'ce':
criterion = celoss
elif opt.loss == 'iou':
criterion = mask_iou_loss
elif opt.loss == 'both':
criterion = lambda pred, target, obj: celoss(
pred, target, obj) + mask_iou_loss(pred, target, obj)
else:
raise TypeError('unknown training loss %s' % opt.loss)
optimizer = None
if opt.solver == 'sgd':
optimizer = optim.SGD(net.parameters(),
lr=opt.learning_rate,
momentum=opt.momentum[0],
weight_decay=opt.weight_decay)
elif opt.solver == 'adam':
optimizer = optim.Adam(net.parameters(),
lr=opt.learning_rate,
betas=opt.momentum,
weight_decay=opt.weight_decay)
else:
raise TypeError('unkown solver type %s' % opt.solver)
# Resume
title = 'STAN'
minloss = float('inf')
opt.checkpoint = osp.join(osp.join(opt.checkpoint, opt.valset))
if not osp.exists(opt.checkpoint):
os.mkdir(opt.checkpoint)
if opt.resume:
# Load checkpoint.
print('==> Resuming from checkpoint {}'.format(opt.resume))
assert os.path.isfile(
opt.resume), 'Error: no checkpoint directory found!'
# opt.checkpoint = os.path.dirname(opt.resume)
checkpoint = torch.load(opt.resume)
minloss = checkpoint['minloss']
start_epoch = checkpoint['epoch']
net.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
skips = checkpoint['max_skip']
try:
if isinstance(skips, list):
for idx, skip in enumerate(skips):
trainloader.dataset.datasets[idx].set_max_skip(skip)
else:
trainloader.dataset.set_max_skip(skip)
except:
print('[Warning] Initializing max skip fail')
logger = Logger(os.path.join(opt.checkpoint, opt.mode + '_log.txt'),
resume=True)
else:
if opt.initial:
print('==> Initialize model with weight file {}'.format(
opt.initial))
weight = torch.load(opt.initial)
if isinstance(weight, OrderedDict):
net.load_param(weight)
else:
net.load_param(weight['state_dict'])
logger = Logger(os.path.join(opt.checkpoint, opt.mode + '_log.txt'),
resume=False)
start_epoch = 0
logger.set_items(['Epoch', 'LR', 'Train Loss'])
# Train and val
for epoch in range(start_epoch):
adjust_learning_rate(optimizer, epoch, opt)
for epoch in range(start_epoch, opt.epochs):
print('\nEpoch: [%d | %d] LR: %f' %
(epoch + 1, opt.epochs, opt.learning_rate))
adjust_learning_rate(optimizer, epoch, opt)
train_loss = train(trainloader,
model=net,
criterion=criterion,
optimizer=optimizer,
epoch=epoch,
use_cuda=use_gpu,
iter_size=opt.iter_size,
mode=opt.mode,
threshold=opt.iou_threshold)
if (epoch + 1) % opt.epoch_per_test == 0:
test_loss = test(testloader,
model=net,
criterion=criterion,
epoch=epoch,
use_cuda=use_gpu,
opt=opt)
# append logger file
logger.log(epoch + 1, opt.learning_rate, train_loss)
# adjust max skip
if (epoch + 1) % opt.epochs_per_increment == 0:
if isinstance(trainloader.dataset, data.ConcatDataset):
for dataset in trainloader.dataset.datasets:
dataset.increase_max_skip()
else:
trainloader.dataset.increase_max_skip()
# save model
is_best = train_loss <= minloss
minloss = min(minloss, train_loss)
skips = [ds.max_skip for ds in trainloader.dataset.datasets] \
if isinstance(trainloader.dataset, data.ConcatDataset) \
else trainloader.dataset.max_skip
save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': net.state_dict(),
'loss': train_loss,
'minloss': minloss,
'optimizer': optimizer.state_dict(),
'max_skip': skips,
},
epoch + 1,
is_best,
checkpoint=opt.checkpoint,
filename=opt.mode)
logger.close()
print('minimum loss:')
print(minloss)
def main():
settings_print_interval = 1 # How often to print loss and other info
settings_batch_size = 4 # Batch size 80 default 64
settings_num_workers = 16 # Number of workers for image loading
settings_normalize_mean = [0.485, 0.456, 0.406] # Normalize mean (default pytorch ImageNet values)
settings_normalize_std = [0.229, 0.224, 0.225] # Normalize std (default pytorch ImageNet values)
settings_search_area_factor = 4.0 # Image patch size relative to target size
settings_feature_sz = 24 # Size of feature map
settings_output_sz = settings_feature_sz * 16 # Size of input image patches 24*16
settings_segm_use_distance = True
# Settings for the image sample and proposal generation
settings_center_jitter_factor = {'train': 0, 'test1': 1.5, 'test2': 1.5}
settings_scale_jitter_factor = {'train': 0, 'test1': 0.25, 'test2': 0.25}
####################################################################################################
start_epoch = 0
random.seed(0)
args = parse_args()
# Use GPU
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu if args.gpu != '' else str(opt.gpu_id)
use_gpu = torch.cuda.is_available() and (args.gpu != '' or int(opt.gpu_id)) >= 0
gpu_ids = [int(val) for val in args.gpu.split(',')]
if not os.path.isdir(opt.checkpoint):
os.makedirs(opt.checkpoint)
# Data
print('==> Preparing dataset')
input_size = opt.input_size
train_transformer = TrainTransform(size=input_size)
#train_transformer = TrainTransform_Noresize()
test_transformer = TestTransform(size=input_size)
try:
if isinstance(opt.trainset, list):
datalist = []
for dataset, freq, max_skip in zip(opt.trainset, opt.datafreq, opt.max_skip):
ds = DATA_CONTAINER[dataset](
train=True,
sampled_frames=opt.sampled_frames,
transform=train_transformer,
max_skip=max_skip,
samples_per_video=opt.samples_per_video
)
datalist += [ds] * freq
trainset = data.ConcatDataset(datalist)
else:
max_skip = opt.max_skip[0] if isinstance(opt.max_skip, list) else opt.max_skip
trainset = DATA_CONTAINER[opt.trainset](
train=True,
sampled_frames=opt.sampled_frames,
transform=train_transformer,
max_skip=max_skip,
samples_per_video=opt.samples_per_video
)
except KeyError as ke:
print('[ERROR] invalide dataset name is encountered. The current acceptable datasets are:')
print(list(DATA_CONTAINER.keys()))
exit()
testset = DATA_CONTAINER[opt.valset](
train=False,
transform=test_transformer,
samples_per_video=1
)
trainloader = data.DataLoader(trainset, batch_size=opt.train_batch, shuffle=True, num_workers=opt.workers,
collate_fn=multibatch_collate_fn, drop_last=True)
testloader = data.DataLoader(testset, batch_size=1, shuffle=False, num_workers=opt.workers,
collate_fn=multibatch_collate_fn)
#########################################################################################
vos_train = Vos(split='train')
transform_train = torchvision.transforms.Compose([dltransforms.ToTensorAndJitter(0.2),
torchvision.transforms.Normalize(mean=settings_normalize_mean,
std=settings_normalize_std)])
data_processing_train = segm_processing.SegmProcessing(search_area_factor=settings_search_area_factor,
output_sz=settings_output_sz,
center_jitter_factor=settings_center_jitter_factor,
scale_jitter_factor=settings_scale_jitter_factor,
mode='pair',
transform=transform_train,
use_distance=settings_segm_use_distance)
dataset_train = segm_sampler.SegmSampler([vos_train], [1],
samples_per_epoch=1000 * settings_batch_size * 8, max_gap=50,
processing=data_processing_train)
loader_train = LTRLoader('train', dataset_train, training=True, batch_size=settings_batch_size,
num_workers=settings_num_workers,
shuffle=True, drop_last=True, stack_dim=1)
#########################################################################################
# Model
print("==> creating model")
net = AMB(opt.keydim, opt.valdim, 'train', mode=opt.mode, iou_threshold=opt.iou_threshold)
print(' Total params: %.2fM' % (sum(p.numel() for p in net.parameters())/1000000.0))
net.eval()
if use_gpu:
net = net.cuda()
assert opt.train_batch % len(gpu_ids) == 0
net = nn.DataParallel(net, device_ids=gpu_ids, dim=0)
# set training parameters
#for p in net.parameters():
# p.requires_grad = True
for name, param in net.named_parameters():
#print(name)
if 'Encoder' in name:
param.requires_grad = False # 冻结 backbone 梯度
else:
param.requires_grad = True
criterion = None
celoss = cross_entropy_loss
if opt.loss == 'ce':
criterion = celoss
elif opt.loss == 'iou':
criterion = mask_iou_loss
elif opt.loss == 'both':
criterion = lambda pred, target, obj: celoss(pred, target, obj) + mask_iou_loss(pred, target, obj)
else:
raise TypeError('unknown training loss %s' % opt.loss)
optimizer = None
if opt.solver == 'sgd':
optimizer = optim.SGD(net.parameters(), lr=opt.learning_rate,
momentum=opt.momentum[0], weight_decay=opt.weight_decay)
elif opt.solver == 'adam':
optimizer = optim.Adam(net.parameters(), lr=opt.learning_rate,
betas=opt.momentum, weight_decay=opt.weight_decay)
else:
raise TypeError('unkown solver type %s' % opt.solver)
# Resume
title = 'Appearance Memory Bank'
minloss = float('inf')
opt.checkpoint = osp.join(osp.join(opt.checkpoint, opt.valset))
if not osp.exists(opt.checkpoint):
os.mkdir(opt.checkpoint)
if opt.resume:
# Load checkpoint.
print('==> Resuming from checkpoint {}'.format(opt.resume))
assert os.path.isfile(opt.resume), 'Error: no checkpoint directory found!'
# opt.checkpoint = os.path.dirname(opt.resume)
checkpoint = torch.load(opt.resume)
minloss = checkpoint['minloss']
start_epoch = checkpoint['epoch']
net.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
skips = checkpoint['max_skip']
try:
if isinstance(skips, list):
for idx, skip in enumerate(skips):
trainloader.dataset.datasets[idx].set_max_skip(skip)
else:
trainloader.dataset.set_max_skip(skip)
except:
print('[Warning] Initializing max skip fail')
logger = Logger(os.path.join(opt.checkpoint, opt.mode+'_log.txt'), resume=True)
else:
if opt.initial:
print('==> Initialize model with weight file {}'.format(opt.initial))
weight = torch.load(opt.initial)
if isinstance(weight, OrderedDict):
net.module.load_param(weight)
else:
net.module.load_param(weight['state_dict'])
logger = Logger(os.path.join(opt.checkpoint, opt.mode+'_log.txt'), resume=False)
start_epoch = 0
logger.set_items(['Epoch', 'LR', 'Train Loss'])
# Train and val
for epoch in range(start_epoch):
adjust_learning_rate(optimizer, epoch, opt)
for epoch in range(start_epoch, opt.epochs):
print('\nEpoch: [%d | %d] LR: %f' % (epoch + 1, opt.epochs, opt.learning_rate))
adjust_learning_rate(optimizer, epoch, opt)
net.module.phase = 'train'
train_loss = train(loader_train, # loader_train trainloader
model=net,
criterion=criterion,
optimizer=optimizer,
epoch=epoch,
use_cuda=use_gpu,
iter_size=opt.iter_size,
mode=opt.mode,
threshold=opt.iou_threshold)
if (epoch + 1) % opt.epoch_per_test == 0:
net.module.phase = 'test'
test_loss = test(testloader,
model=net.module,
criterion=criterion,
epoch=epoch,
use_cuda=use_gpu)
# append logger file
logger.log(epoch+1, opt.learning_rate, train_loss)
# adjust max skip
if (epoch + 1) % opt.epochs_per_increment == 0:
if isinstance(trainloader.dataset, data.ConcatDataset):
for dataset in trainloader.dataset.datasets:
dataset.increase_max_skip()
else:
trainloader.dataset.increase_max_skip()
# save model
is_best = train_loss <= minloss
minloss = min(minloss, train_loss)
skips = [ds.max_skip for ds in trainloader.dataset.datasets] \
if isinstance(trainloader.dataset, data.ConcatDataset) \
else trainloader.dataset.max_skip
save_checkpoint({
'epoch': epoch + 1,
'state_dict': net.state_dict(),
'loss': train_loss,
'minloss': minloss,
'optimizer': optimizer.state_dict(),
'max_skip': skips,
}, epoch + 1, is_best, checkpoint=opt.checkpoint, filename=opt.mode)
logger.close()
print('minimum loss:')
print(minloss)