python anotation_Python data.AnnotationTransform方法代码示例

# 需要导入模块: import data [as 别名]

# 或者: from data import AnnotationTransform [as 别名]

def train():

net.train()

epoch = 0 + args.resume_epoch

print('Loading Dataset...')

dataset = VOCDetection(args.training_dataset, preproc_s3fd(img_dim, rgb_means, cfg['max_expand_ratio']), AnnotationTransform())

epoch_size = math.ceil(len(dataset) / args.batch_size)

max_iter = args.max_epoch * epoch_size

stepvalues = (200 * epoch_size, 250 * epoch_size)

step_index = 0

if args.resume_epoch > 0:

start_iter = args.resume_epoch * epoch_size

else:

start_iter = 0

for iteration in range(start_iter, max_iter):

if iteration % epoch_size == 0:

# create batch iterator

batch_iterator = iter(data.DataLoader(dataset, batch_size, shuffle=True, num_workers=args.num_workers, collate_fn=detection_collate, pin_memory=True))

if (epoch % 10 == 0 and epoch > 0) or (epoch % 5 == 0 and epoch > 200):

torch.save(net.state_dict(), args.save_folder + 'S3FD_epoch_' + repr(epoch) + '.pth')

epoch += 1

load_t0 = time.time()

if iteration in stepvalues:

step_index += 1

lr = adjust_learning_rate(optimizer, args.gamma, epoch, step_index, iteration, epoch_size)

# load train data

images, targets = next(batch_iterator)

if args.cuda:

images = Variable(images.cuda())

targets = [Variable(anno.cuda()) for anno in targets]

else:

images = Variable(images)

targets = [Variable(anno) for anno in targets]

# forward

out = net(images)

# backprop

optimizer.zero_grad()

loss_l, loss_c = criterion(out, priors, targets)

loss = loss_l + cfg['conf_weight'] * loss_c

loss.backward()

optimizer.step()

load_t1 = time.time()

print('Epoch:' + repr(epoch) + ' || epochiter: ' + repr(iteration % epoch_size) + '/' + repr(epoch_size) +

'|| Totel iter ' + repr(iteration) + ' || L: %.4f C: %.4f||' % (loss_l.item(), cfg['conf_weight'] * loss_c.item()) +

'Batch time: %.4f sec. ||' % (load_t1 - load_t0) + 'LR: %.8f' % (lr))

if writer is not None:

writer.add_scalar('train/loss_l', loss_l.item(), iteration)

writer.add_scalar('train/loss_c', cfg['conf_weight'] * loss_c.item(), iteration)

writer.add_scalar('train/lr', lr, iteration)

torch.save(net.state_dict(), args.save_folder + 'Final_S3FD.pth')