def main(args):
    torch.backends.cudnn.benchmark = True
    device = torch.device(('cuda' if (args.device == 'cuda') else 'cpu'))
    (train_data, val_data) = load_data(args.train_path, args.val_path, extensions=args.extensions)
    train_loader = DataLoader(train_data, batch_size=args.batch_size, shuffle=True)
    val_loader = DataLoader(val_data, batch_size=1, shuffle=False)
    model = FC_Siam_Diff(num_classes=args.num_classes)
    model.to(device)
    if args.resume:
        model.load_state_dict(torch.load(args.resume, map_location=device))
    criterion = nn.CrossEntropyLoss().to(device)
    criterion = nn.BCELoss()
    optimizer = optim.Adam(model.parameters(), lr=args.lr)
    lr_scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=10, T_mult=1, eta_min=1e-08)
    writer = SummaryWriter(args.ckp_dir)
    for epoch in range(args.epochs):
        writer.add_scalar('train/lr', lr_scheduler.get_lr()[0], epoch)
        train_one_epoch(epoch, train_loader, model, criterion, optimizer, device, writer)
        evalidation(epoch, val_loader, model, criterion, device, writer, args.tb_test_imgs)
        lr_scheduler.step()
        if ((epoch % 2) == 0):
            torch.save(model.state_dict(), os.path.join(args.ckp_dir, 'cd_epoch_{}.pth'.format(epoch)))
def main(args):
    torch.backends.cudnn.benchmark = True
    device = torch.device(('cuda' if (args.device == 'cuda') else 'cpu'))
    (train_data, val_data) = load_data(args.train_path, args.val_path, extensions=args.extensions)
    train_loader = DataLoader(train_data, batch_size=args.batch_size, shuffle=True)
    val_loader = DataLoader(val_data, batch_size=1, shuffle=False)
    model = FC_Siam_Diff(num_classes=args.num_classes)
    model.to(device)
    if args.resume:
        model.load_state_dict(torch.load(args.resume, map_location=device))
    criterion = nn.CrossEntropyLoss().to(device)
    criterion = nn.BCELoss()
    optimizer = optim.Adam(model.parameters(), lr=args.lr)
    lr_scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=10, T_mult=1, eta_min=1e-08)
    writer = SummaryWriter(args.ckp_dir)
    for epoch in range(args.epochs):
        writer.add_scalar('train/lr', lr_scheduler.get_lr()[0], epoch)
        train_one_epoch(epoch, train_loader, model, criterion, optimizer, device, writer)
        evalidation(epoch, val_loader, model, criterion, device, writer, args.tb_test_imgs)
        lr_scheduler.step()
        if ((epoch % 2) == 0):
            torch.save(model.state_dict(), os.path.join(args.ckp_dir, 'cd_epoch_{}.pth'.format(epoch)))