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))