<jupyter_start><jupyter_text>Train a BBNet in rounds<jupyter_code>import dorsalventral.models.targets as targets
import dorsalventral.vis_utils as vis_utils
import dorsalventral.models.fire_propagation as fprop
import dorsalventral.models.layer_utils as layer_utils
tdw_dataset = raft_datasets.TdwFlowDataset(
    root='/data5/dbear/tdw_datasets/playroom_large_v3copy/',
#     root='/mnt/fs5/dbear/tdw_datasets/playroom_large_v3/',
    dataset_names=['model_split_4'],    
#     dataset_names=['model_split_%d' % split for split in range(4)],
    split='training',
    filepattern="00*1",
#     test_filepattern="0*[0-4]",
    test_filepattern="*9",
    min_start_frame=5,
    max_start_frame=5,
    scale_to_pixels=True,
    get_gt_flow=True,
    get_backward_frame=True
)
print(len(tdw_dataset))
print(tdw_dataset.is_test)<jupyter_output>10
False<jupyter_text>Load the motion model for round 0<jupyter_code>motion_path = '../checkpoints/35000_motion-rnd0-tdw-bs8-small-fullplay-tr0-8.pth'
mclass_net = load_model(
                       motion_path,
                       small=True, 
                       cuda=True, 
                       train=False)
mclass_net.module.freeze_bn()
mclass_net.eval()
print(layer_utils.num_parameters(mclass_net))
# ex = 0
# viz_raft_prediction(mclass_net, tdw_dataset, ex=ex, iters=[1,6,12,24], thingness=True, thresh=None)<jupyter_output>/home/dbear/tenv/lib/python3.7/site-packages/ipykernel/ipkernel.py:283: DeprecationWarning: `should_run_async` will not call `transform_cell` automatically in the future. Please pass the result to `transformed_cell` argument and any exception that happen during thetransform in `preprocessing_exc_tuple` in IPython 7.17 and above.
  and should_run_async(code)<jupyter_text>Create a BBNet<jupyter_code>bbnet = fprop.BipartiteBootNet(
    motion_params={'hidden_dim': 64, 'num_iters': 2},
    mode='train_static').cuda()
print("num parameters", layer_utils.num_parameters(bbnet))
def get_data(dataset, ex, up_to=2):
    return [x[None] for x in dataset[ex][:up_to]]

def train_step(data, model, optimizer, **kwargs):
    optimizer.zero_grad()
    preds = model(*[x.cuda() for x in data], **kwargs)
    loss = sum(model.loss_list)
    loss.backward()
    optimizer.step()
    target = getattr(model, 'target', None)
    return (data, preds, target, loss)

def eval_step(data, model, teacher=None, mode='val', out_idx=-1, **kwargs):
    model.mode = mode
    if teacher is not None:
        teacher_motion = teacher(*[x.cuda() for x in data], **kwargs)[out_idx].detach()
        data.append(teacher_motion)
    preds = model(*[x.cuda() for x in data], **kwargs)
    return (data, preds)

def train_with_teacher(dataloader, model, teacher, mode='train_static',
                       lr=0.0003, num_steps=100, print_every=5, out_idx=-1,
                       teacher_iters=1, 
                       load_path=None, save_path=None,
                       **kwargs):
    
    if load_path is not None:
        did_load = model.load_state_dict(torch.load(load_path))
        print("loaded %s, %s" % (load_path, did_load))
        
    model.train()
    model.mode = mode
    if mode == 'train_static':
        optimizer = torch.optim.Adam(model.static_model.parameters(), lr=lr)
    elif mode == 'train_motion':
        optimizer = torch.optim.Adam(model.motion_model.parameters(), lr=lr)
    loss_curve = []
    step, epoch, running_loss = 0, 0, 0.0
    while (step <= num_steps):
        epoch += 1
        for _step, data in enumerate(dataloader):
            data = data[:2]
            teacher_motion = eval_step(data, teacher, test_mode=True, iters=teacher_iters)[-1][out_idx].detach()
            data.append(teacher_motion)
            loss = train_step(data, model, optimizer, **kwargs)[-1]
            loss_curve.append(loss.item())
            running_loss += (loss.item()) / print_every
        
            if (step % print_every) == 0:
                print("%s w teacher epoch %d / step %d ---- loss %.3f" % (model.mode, epoch-1, _step, running_loss))
                running_loss = 0.0
                
            step += 1
            
    if save_path is not None:
        torch.save(model.state_dict(), save_path)
        print("saved model weights to %s" % save_path)
        
    return loss_curve

def train_selfsup(dataloader, model, mode='train_motion',
                  num_steps=100, print_every=5, lr=0.0003, 
                  load_path=None, save_path=None,
                  **kwargs):
    
    if load_path is not None:
        did_load = model.load_state_dict(torch.load(load_path))
        print("loaded %s, %s" % (load_path, did_load))    
    
    model.train()
    assert mode in ['train_static', 'train_motion', 'train_both'], mode
    model.mode = mode
    if mode == 'train_motion':
        params = model.motion_model.parameters()
    elif mode == 'train_static':
        params = model.static_model.parameters()
    elif mode == 'train_both':
        params = model.parameters()
    optimizer = torch.optim.Adam(params, lr=lr)
    loss_curve = []
    step, epoch, running_loss = 0, 0, 0.0
    while (step <= num_steps):
        epoch += 1
        for _step, data in enumerate(dataloader):
            data = data[:2]
            loss = train_step(data, model, optimizer, **kwargs)[-1]
            loss_curve.append(loss.item())
            running_loss += (loss.item()) / print_every
        
            if (step % print_every) == 0:
                print("%s epoch %d / step %d ---- loss %.3f" % (model.mode, epoch-1, _step, running_loss))
                running_loss = 0.0
                
            step += 1
            
    if save_path is not None:
        torch.save(model.state_dict(), save_path)
        print("saved model weights to %s" % save_path)
            
    return loss_curve    
bs = 1
dataloader = torch.utils.data.DataLoader(tdw_dataset, batch_size=bs, shuffle=True)


# load_path = '../checkpoints/bbnet_split4ex1_demo.pth'
# did_load = bbnet.load_state_dict(torch.load(load_path))
# print(did_load)

# bbnet_rnd0 = fprop.BipartiteBootNet(
#     motion_params={'hidden_dim': 64, 'num_iters': 2},
#     mode='train_static').cuda()
# bbnet_rnd0.eval()
# bbnet_rnd0.mode = 'val'
# did_load = bbnet_rnd0.load_state_dict(torch.load(load_path))
# print(did_load)


# static_loss_curve = train_with_teacher(dataloader, bbnet, mclass_net, mode='train_static', num_steps=1000)
# motion_loss_curve = train_with_teacher(dataloader, bbnet, mclass_net, mode='train_motion', 
#                                        load_path='../checkpoints/bbnet_split4ex10_rnd0static.pth',
#                                        save_path='../checkpoints/bbnet_split4ex10_rnd0motion.pth',                                       
#                                        num_steps=1000)
# selfsup_loss_curve = train_selfsup(dataloader, bbnet, mode='train_static', 
#                                    save_path='../checkpoints/bbnet_split4ex10_rnd1static.pth',
#                                    lr=0.0001, 
#                                    num_steps=1000)

# selfsup_loss_curve_motion = train_with_teacher(dataloader, bbnet, bbnet_rnd0, 
#                                                out_idx=0, mode='train_motion', 
#                                                num_steps=100)
# selfsup_loss_curve_static = train_selfsup(dataloader, bbnet, mode='train_static', lr=0.0001, num_steps=100)


# selfsup_loss_curve_motion = train_selfsup(dataloader, bbnet, mode='train_motion', lr=1e-6, num_steps=100)

# sv_path = '../checkpoints/bbnet_split4ex10_rnd0static.pth'
# bbnet.load_state_dict(torch.load(sv_path))
ex = 1
data, preds = eval_step(get_data(tdw_dataset, ex), bbnet, teacher=None, 
                        iters=1, test_mode=True,
                        mode='train_both')
fig, axes = plt.subplots(1,5,figsize=(20,4))
vis_utils.imshow(data[0]/255., axes[0])
titles = ['motion', 'object prior', 'static target', 'motion_target']
axes[0].set_title('image')
axes[0].set_xticks([])
axes[0].set_yticks([])
for i in range(1,5):
    plot = preds[i-1]
    K = plot.shape[1]
    if K != 1:
        plot = plot[:,(K//2),None]
        
    if i in [1,2]:
        plot = torch.sigmoid(plot)
        vis_utils.imshow(plot, axes[i], cmap='magma', vmin=0, vmax=1)
    if i in [3,4]:
        vis_utils.imshow(plot, axes[i], cmap='RdBu_r', vmin=-1, vmax=1)
    axes[i].set_title(titles[i-1], fontsize=20)
    axes[i].set_xticks([])
    axes[i].set_yticks([])
plt.tight_layout()
plt.show()
prop = fprop.MotionPropagator(
    num_samples=3, 
    confidence_thresh=3,
    num_initial_points=None,
    motion_thresh=0.5,
    radius=1, 
    num_iters=100, 
    update_thresh=0.25).cuda()
h, fire = prop(preds[0], preds[1])
vis_utils.imshow(h)
vis_utils.imshow(fire, cmap='inferno', vmin=0, vmax=fire.amax().item())<jupyter_output>/home/dbear/tenv/lib/python3.7/site-packages/ipykernel/ipkernel.py:283: DeprecationWarning: `should_run_async` will not call `transform_cell` automatically in the future. Please pass the result to `transformed_cell` argument and any exception that happen during thetransform in `preprocessing_exc_tuple` in IPython 7.17 and above.
  and should_run_async(code)
{'fraction_non_alphanumeric': 0.12748614281056406, 'fraction_numerical': 0.02271492229105532, 'mean_word_length': 2.8071990070335127, 'pattern_counts': {'":': 0, '<': 12, '<?xml version=': 0, '>': 10, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 11, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '41019', 'license': 'BSD-3-Clause', 'n_tokens_mistral': 3311, 'n_tokens_neox': 3122, 'n_words': 584, 'path': 'notebooks/.ipynb_checkpoints/BBNet_db032622-checkpoint.ipynb', 'repo_name': 'neuroailab/RAFT', 'size': '9201'}