<jupyter_start><jupyter_text>Load the pts files`scan_line_tensor` is the data file, a 3-D tensor of size [num_scan_lines,pts_per_scan_line,num_feats].  `idx_lists` indicate the top left corner of each training or validation square.  `sc` is the minmaxscaler used to generate the training set. Needed here to calculate weightedMSELoss at real scale.<jupyter_code># Loads as a list of numpy arrays
scan_line_tensor = torch.load('../lidar_data/old_bad/32_32/'+'scan_line_tensor.pts')
train_idx_list = torch.load('../lidar_data/old_bad/32_32/'+'train_idx_list.pts')
valid_idx_list = torch.load('../lidar_data/old_bad/32_32/'+'valid_idx_list.pts')
sc = torch.load('../lidar_data/old_bad/32_32/'+'sc.pts')
scan_line_tensor = scan_line_tensor[:,:,:3]
def add_mask(sample,mask_pts_per_seq,consecutive=True):
    # Given a 3-D tensor of all ones, returns a mask_tensor of same shape 
    # with random masking determined by mask_pts_per_seq
    mask_tensor = torch.ones(sample.shape)
    seq_len = mask_tensor.shape[0]
    
    if consecutive:
        # Creates a square of missing points
        first_mask = int(np.random.choice(np.arange(8,seq_len-8-mask_pts_per_seq),1))
        
        mask_tensor[first_mask:first_mask+mask_pts_per_seq,first_mask:first_mask+mask_pts_per_seq,0] = 0
            
    else:
        # TO DO: Random points throughout the patch
        for i in range(sample.shape[0]):
            m[i,:] = np.random.choice(np.arange(8,seq_len-8),mask_pts_per_seq,replace=False)

    return mask_tensor
# Dataloader class
class LidarLstmDataset(udata.Dataset):
    def __init__(self, scan_line_tensor, idx_list, seq_len = 64, mask_pts_per_seq = 5, consecutive = True):
        super(LidarLstmDataset, self).__init__()
        self.scan_line_tensor = scan_line_tensor
        self.idx_list = idx_list
        self.seq_len = seq_len
        self.mask_pts_per_seq = mask_pts_per_seq
        self.consecutive = consecutive

    def __len__(self):
        return len(self.idx_list)

    def __getitem__(self,index):
        row = self.idx_list[index][0]
        col = self.idx_list[index][1]
        clean = self.scan_line_tensor[row:row+self.seq_len,col:col+self.seq_len,:]
        mask = add_mask(clean,self.mask_pts_per_seq,self.consecutive)
        return clean.permute(2,0,1), mask.permute(2,0,1)
train_dataset = LidarLstmDataset(scan_line_tensor,train_idx_list,args.seq_len, args.mask_pts_per_seq)
valid_dataset = LidarLstmDataset(scan_line_tensor,valid_idx_list,args.seq_len, args.mask_pts_per_seq)

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32, num_workers=4, shuffle=True)
valid_loader = torch.utils.data.DataLoader(valid_dataset, batch_size=32, num_workers=4, shuffle=True)
# Define a loss function that weights the loss according to coordinate ranges (xmax-xmin, ymax-ymin, zmax-zmin)
def weighted_MSELoss(pred,true,sc,mask_pts_per_seq=args.mask_pts_per_seq):
    '''weighted_MSELoss reconverts MSE loss back to the original scale of x,y,z.
    Rationale is because xyz have such different ranges, we don't want to ignore the ones with largest scale.
    Assumes that x,y,z are the first 3 features in sc scaler'''
    
    ranges = torch.Tensor(sc.data_max_[:3]-sc.data_min_[:3])
    raw_loss = torch.zeros(3,dtype=float)
    for i in range(3):
        raw_loss[i] = F.mse_loss(pred[:,i,:,:], true[:,i,:,:], reduction="sum") 
    return (ranges**2 * raw_loss).sum() #/ (pred.shape[0]*mask_pts_per_seq**2)<jupyter_output><empty_output><jupyter_text>Training the model<jupyter_code># Track moving average of loss values
train_meters = {name: utils.RunningAverageMeter(0.98) for name in (["train_loss"])}
valid_meters = {name: utils.AverageMeter() for name in (["valid_loss"])}
writer = SummaryWriter(log_dir=args.experiment_dir) if not args.no_visual else None
# TRAINING - X ONLY
for epoch in range(start_epoch, args.num_epochs):
    if args.resume_training:
        if epoch %1 == 0:
            optimizer.param_groups[0]["lr"] /= 2
            print('learning rate reduced by factor of 2')

    train_bar = utils.ProgressBar(train_loader, epoch)
    for meter in train_meters.values():
        meter.reset()

#     epoch_loss_sum = 0
    for batch_id, (clean, mask) in enumerate(train_bar):
        # dataloader returns [clean, mask] list
        model.train()
        global_step += 1
        inputs = clean.to(device)
        mask_inputs = mask.to(device)
        # only use the mask part of the outputs
        raw_outputs = torch.zeros(mask_inputs[:,:3,:,:].shape).to(device)
        raw_x = model(inputs,mask_inputs)
        raw_outputs[:,0,:,:] = raw_x[:,0,:,:]
        outputs = (1-mask_inputs[:,:3,:,:])*raw_outputs + mask_inputs[:,:3,:,:]*inputs[:,:3,:,:]
        
        if args.wtd_loss:
            loss = weighted_MSELoss(outputs,inputs[:,:3,:,:],sc)/(inputs.size(0)*(args.mask_pts_per_seq**2))
            # Regularization?
            
        else:
            # normalized by the number of masked points
            loss = F.mse_loss(outputs, inputs[:,:3,:,:], reduction="sum") / \
                   (inputs.size(0) * (args.mask_pts_per_seq**2))

        model.zero_grad()
        loss.backward()
        optimizer.step()
#         epoch_loss_sum += loss * inputs.size(0)
        train_meters["train_loss"].update(loss)
        train_bar.log(dict(**train_meters, lr=optimizer.param_groups[0]["lr"]), verbose=True)

        if writer is not None and global_step % args.log_interval == 0:
            writer.add_scalar("lr", optimizer.param_groups[0]["lr"], global_step)
            writer.add_scalar("loss/train", loss.item(), global_step)
            gradients = torch.cat([p.grad.view(-1) for p in model.parameters() if p.grad is not None], dim=0)
            writer.add_histogram("gradients", gradients, global_step)
            sys.stdout.flush()
#     epoch_loss = epoch_loss_sum / len(train_loader.dataset)
    
    if epoch % args.valid_interval == 0:
        model.eval()
        for meter in valid_meters.values():
            meter.reset()

        valid_bar = utils.ProgressBar(valid_loader)
        val_loss = 0
        for sample_id, (clean, mask) in enumerate(valid_bar):
            with torch.no_grad():
                inputs = clean.to(device)
                mask_inputs = mask.to(device)
                # only use the mask part of the outputs
                raw_output = model(inputs,mask_inputs)
                output = (1-mask_inputs[:,:3,:,:])*raw_output + mask_inputs[:,:3,:,:]*inputs[:,:3,:,:]

                # TO DO, only run loss on masked part of output
                
                if args.wtd_loss:
                    val_loss = weighted_MSELoss(output,inputs[:,:3,:,:],sc)/(inputs.size(0)*(args.mask_pts_per_seq**2))
                else:
                    # normalized by the number of masked points
                    val_loss = F.mse_loss(output, inputs[:,:3,:,:], reduction="sum")/(inputs.size(0)* \
                                                                                    (args.mask_pts_per_seq**2))

                valid_meters["valid_loss"].update(val_loss.item())

        if writer is not None:
            writer.add_scalar("loss/valid", valid_meters['valid_loss'].avg, global_step)
            sys.stdout.flush()

        logging.info(train_bar.print(dict(**train_meters, **valid_meters, lr=optimizer.param_groups[0]["lr"])))
        utils.save_checkpoint(args, global_step, model, optimizer, score=valid_meters["valid_loss"].avg, mode="min")
    scheduler.step()

logging.info(f"Done training! Best Loss {utils.save_checkpoint.best_score:.3f} obtained after step {utils.save_checkpoint.best_step}.")

raw_x.shape
# Single output example
i,m = next(iter(train_loader))
inputs = i.to(device)
mask = m.to(device)
raw_output = model(inputs,mask)
output_model = (1-mask[:,:3,:])*raw_output + mask[:,:3,:]*inputs[:,:3,:]
raw_output.shape<jupyter_output><empty_output><jupyter_text>BaselinesFirst: Interpolate between last and next<jupyter_code>import gc
gc.collect()
# model_cpu = model.to('cpu')
def loss_comparison(loader,model,mask_pts_per_seq=args.mask_pts_per_seq,pt_count=len(valid_dataset)):
    wtd_loss = True
    loss_model = 0
    loss_interp = 0
    for batch_id, (i, m) in enumerate(loader):
        with torch.no_grad():
            # conv1D model
            inputs = i.to(device)
            mask = m.to(device)
            raw_outputs = torch.zeros(mask_inputs[:,:3,:,:].shape).to(device)
            raw_x = model(inputs,mask_inputs)
            raw_outputs[:,0,:,:] = raw_x[:,0,:,:]
            output_model = (1-mask_inputs[:,:3,:,:])*raw_outputs + mask_inputs[:,:3,:,:]*inputs[:,:3,:,:]
            
            # Interpolation
            output_interp = lidar_data_processing.outer_interp_loop(i,m,mask_pts_per_seq,2)

            if wtd_loss:
                loss_model+=weighted_MSELoss(output_model,inputs[:,:3,:],sc)
                loss_interp+=weighted_MSELoss(output_interp,i[:,:3,:],sc)
            else:
                # normalized by the number of masked points
                loss_model += F.mse_loss(output_model, inputs[:,:3,:], reduction="sum") 
                loss_interp += F.mse_loss(output_interp, i[:,:3,:], reduction="sum") 
        print("Batch {} done".format(batch_id))

    # Normalize by number of batches
    loss_model = loss_model/((mask_pts_per_seq**2)*pt_count)
    loss_interp = loss_interp/((mask_pts_per_seq**2)*pt_count)
    print("Validation Loss\n","*"*30)
    print("Model: {:2.2f}".format(loss_model))
    print("Interpolation: {:2.2f}".format(loss_interp))
    gc.collect()
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=128, num_workers=8, shuffle=True)
# loss_comparison(train_loader,model)
loss_comparison(train_loader,model,pt_count=len(train_dataset))
valid_loader = torch.utils.data.DataLoader(valid_dataset, batch_size=128, num_workers=8, shuffle=True)
# loss_comparison(train_loader,model)
loss_comparison(valid_loader,model,pt_count=len(valid_dataset))<jupyter_output><empty_output><jupyter_text>Plotting<jupyter_code>inputs,mask_inputs = next(iter(valid_loader))
# only use the mask part of the outputs

raw_outputs = torch.zeros(mask_inputs[:,:3,:,:].shape).to(device)
raw_x = model(inputs.to(device),mask_inputs.to(device))
raw_outputs[:,0,:,:] = raw_x[:,0,:,:]
outputs = (1-mask_inputs[:,:3,:,:])*raw_outputs.to('cpu') + mask_inputs[:,:3,:,:]*inputs[:,:3,:,:]

output_interp = lidar_data_processing.outer_interp_loop(inputs,mask_inputs,args.mask_pts_per_seq,2)
%matplotlib notebook
im_idx = 28

def plot_infill(clean,output,mask,label_key = 'z'):
    ''' Plotting function for 2D infill. Takes a single patch for clean,output, 
    and mask, and a label_key indicating which value (x,y, or z) to display.    
    '''
    # Which dimension to plot
    xyz_dict = {'x':2,'y':3,'z':4}
    z_val = xyz_dict[label_key]
    # Set up plot
    fig = plt.figure(figsize=[12,12])
    ax = fig.add_subplot(111, projection='3d')

    # Plot unmasked points
    surrounding_no_mask = mask[0] != 0
    unmasked_arr = np.array(lidar_data_processing.surrounding_grid(clean,surrounding_no_mask))
    ax.scatter(unmasked_arr[:,0],unmasked_arr[:,1],unmasked_arr[:,z_val], marker='+')

    # Plot masked, filled points
    surrounding_mask = mask[0] == 0
    filled_arr = np.array(lidar_data_processing.surrounding_grid(output,surrounding_mask))
    ax.scatter(filled_arr[:,0],filled_arr[:,1],filled_arr[:,z_val], color='r', marker='o')

    # Plot original, masked points
    masked_arr = np.array(lidar_data_processing.surrounding_grid(clean,surrounding_mask))
    ax.scatter(masked_arr[:,0],masked_arr[:,1],masked_arr[:,z_val], color='g', marker='o')

    # Labels and such
    ax.set_xlabel('Grid Across-Flight',fontsize=15)
    ax.set_ylabel('Grid Along-Flight',fontsize=15)
    ax.set_zlabel(label_key+' value',fontsize=15)

    plt.show()
#     print(output.shape)
#     print(clean.shape)
#     print(weighted_MSELoss(output.unsqueeze(0),clean[:,:3,:].unsqueeze(0),sc))
          
plot_infill(inputs[im_idx],outputs[im_idx],mask_inputs[im_idx],'z')
plot_infill(inputs[im_idx],output_interp[im_idx],mask_inputs[im_idx],'z')
import plotly
import plotly.graph_objs as go
outputs.shape
# Configure Plotly to be rendered inline in the notebook.
plotly.offline.init_notebook_mode()

# Configure the trace.
trace = go.Scatter3d(
    x=outputs.detach()[0,0,:,5],  # <-- Put your data instead
    y=outputs.detach()[0,1,:,5],  # <-- Put your data instead
    z=outputs.detach()[0,2,:,5],  # <-- Put your data instead
    mode='markers',
    marker={
        'size': 10,
        'opacity': 0.8,
    }
)

# Configure the layout.
layout = go.Layout(
    margin={'l': 0, 'r': 0, 'b': 0, 't': 0}
)

data = [trace]

plot_figure = go.Figure(data=data, layout=layout)

# Render the plot.
plotly.offline.iplot(plot_figure)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13987701408889236, 'fraction_numerical': 0.01478944500661633, 'mean_word_length': 3.623245771860381, 'pattern_counts': {'":': 0, '<': 20, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '340757', 'license': 'MIT', 'n_tokens_mistral': 4581, 'n_tokens_neox': 4483, 'n_words': 790, 'path': 'lidar_conv2d-x-only.ipynb', 'repo_name': 'AccomplishedCode/point_cloud_experiments', 'size': '12847'}