def BMN_Train(opt):
writer = SummaryWriter()
model = BMN(opt).cuda()
optimizer = optim.Adam(model.parameters(),
lr=opt["training_lr"],
weight_decay=opt["weight_decay"])
train_loader = torch.utils.data.DataLoader(VideoDataSet(opt,
subset="train"),
batch_size=opt["batch_size"],
shuffle=True,
num_workers=8,
pin_memory=True)
test_loader = torch.utils.data.DataLoader(VideoDataSet(
opt, subset="validation"),
batch_size=opt["batch_size"],
shuffle=False,
num_workers=8,
pin_memory=True)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=opt["step_size"],
gamma=opt["step_gamma"])
for epoch in range(opt["train_epochs"]):
scheduler.step()
train_BMN(train_loader, model, optimizer, epoch, writer, opt)
test_BMN(test_loader, model, epoch, writer, opt)
writer.close()
def BMN_Train(opt):
model = BMN(opt)
model = torch.nn.DataParallel(model, device_ids=[0, 1]).cuda()
optimizer = optim.Adam(filter(lambda p: p.requires_grad,
model.parameters()),
lr=opt["training_lr"],
weight_decay=opt["weight_decay"])
train_loader = torch.utils.data.DataLoader(VideoDataSet(opt,
subset="train"),
batch_size=opt["batch_size"],
shuffle=True,
num_workers=8,
pin_memory=True)
test_loader = torch.utils.data.DataLoader(VideoDataSet(
opt, subset="validation"),
batch_size=opt["batch_size"],
shuffle=False,
num_workers=8,
pin_memory=True)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=opt["step_size"],
gamma=opt["step_gamma"])
bm_mask = get_mask(opt["temporal_scale"])
for epoch in range(opt["train_epochs"]):
scheduler.step()
train_BMN(train_loader, model, optimizer, epoch, bm_mask)
test_BMN(test_loader, model, epoch, bm_mask)
def BMN_Train(opt):
start_time = time.time()
model = BMN(opt)
model = torch.nn.DataParallel(model,
device_ids=list(range(opt['n_gpu']))).cuda()
print('using {} gpus to train!'.format(opt['n_gpu']))
optimizer = optim.Adam(filter(lambda p: p.requires_grad,
model.parameters()),
lr=opt["training_lr"],
weight_decay=opt["weight_decay"])
train_loader = torch.utils.data.DataLoader(VideoDataSet(opt,
subset="train"),
batch_size=opt["batch_size"],
shuffle=True,
num_workers=opt['num_workers'],
pin_memory=True)
test_loader = torch.utils.data.DataLoader(VideoDataSet(
opt, subset="validation"),
batch_size=opt["batch_size"],
shuffle=False,
num_workers=opt['num_workers'],
pin_memory=True)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=opt["step_size"],
gamma=opt["step_gamma"])
bm_mask = get_mask(opt["temporal_scale"])
best_loss = 1e10
for epoch in range(opt["train_epochs"]):
train_BMN(train_loader, model, optimizer, epoch, bm_mask)
best_loss = test_BMN(test_loader, model, epoch, bm_mask, best_loss)
scheduler.step()
print("Total time (BMN_Train):",
datetime.timedelta(seconds=time.time() - start_time))