Function train

make_decoder.py:25–55

(args)

    return (output.float().cpu(), target.cpu()) if iscpu else (output.float(), target)

def train(args):
    print (f'training decoder with N_L={args.msg_len}')
    dataset = MsgDataset(msg_len=args.msg_len, max_size=args.max_size)
    trainloader = DataLoader(dataset, batch_size=min(args.batch_size, dataset.__len__()), shuffle=True)
    testloader = DataLoader(dataset, batch_size=min(args.batch_size, dataset.__len__()), shuffle=False)
    model = CLIPWatermarker(msg_len=args.msg_len).to(DEVICE)
    optimizer = optim.Adam(model.msg_decoder.parameters(), lr=args.lr)
    loss_fn = nn.BCEWithLogitsLoss()
    all_epochs = trange(1, args.num_epochs + 1)

    for eIdx in all_epochs:
        loss_value = 0
        for batch in trainloader:
            optimizer.zero_grad()
            loss = loss_fn(*single_step_inference(batch, model))
            loss.backward()
            optimizer.step()
            loss_value += loss.item()

        score = test(args, model, testloader)
        model.train()

        all_epochs.set_description(f'BAcc : {score} | Loss : {loss_value / dataset.__len__()}')

    if args.save:
        os.makedirs(args.sdir, exist_ok=True)
        sdict = {
            'model'   : model.eval().cpu().state_dict(),
            **{key : getattr(dataset, key) for key in ['data', 'b2t_maps']}
        }
        torch.save(sdict, os.path.join(args.sdir, f'CLIP-MsgDecoder-{args.msg_len}.pkl'))

def test(args, model=None, testloader=None):
    if model is None: