<jupyter_start><jupyter_text>Model Definition<jupyter_code>class LyftMultiModel(nn.Module):

    def __init__(self, cfg: Dict, num_modes=3):
        super().__init__()

        architecture = cfg["model_params"]["model_architecture"]
        backbone = eval(architecture)(pretrained=True, progress=True)
        self.backbone = backbone

        num_history_channels = (cfg["model_params"]["history_num_frames"] + 1) * 2
        num_in_channels = 3 + num_history_channels

        self.backbone.conv1 = nn.Conv2d(
            num_in_channels,
            self.backbone.conv1.out_channels,
            kernel_size=self.backbone.conv1.kernel_size,
            stride=self.backbone.conv1.stride,
            padding=self.backbone.conv1.padding,
            bias=False,
        )

        if architecture == "resnet50":
            backbone_out_features = 2048
        else:
            backbone_out_features = 512

        # X, Y coords for the future positions (output shape: batch_sizex50x2)
        self.future_len = cfg["model_params"]["future_num_frames"]
        num_targets = 2 * self.future_len

        # You can add more layers here.
        self.head = nn.Sequential(
            # nn.Dropout(0.2),
            nn.Linear(in_features=backbone_out_features, out_features=4096),
        )

        self.num_preds = num_targets * num_modes
        self.num_modes = num_modes

        self.logit = nn.Linear(4096, out_features=self.num_preds + num_modes)

    def forward(self, x):
        x = self.backbone.conv1(x)
        x = self.backbone.bn1(x)
        x = self.backbone.relu(x)
        x = self.backbone.maxpool(x)

        x = self.backbone.layer1(x)
        x = self.backbone.layer2(x)
        x = self.backbone.layer3(x)
        x = self.backbone.layer4(x)

        x = self.backbone.avgpool(x)
        x = torch.flatten(x, 1)

        x = self.head(x)
        x = self.logit(x)

        # pred (batch_size)x(modes)x(time)x(2D coords)
        # confidences (batch_size)x(modes)
        bs, _ = x.shape
        pred, confidences = torch.split(x, self.num_preds, dim=1)
        pred = pred.view(bs, self.num_modes, self.future_len, 2)
        assert confidences.shape == (bs, self.num_modes)
        confidences = torch.softmax(confidences, dim=1)
        return pred, confidences

def forward(data, model, device, criterion=pytorch_neg_multi_log_likelihood_batch, compute_loss=True):
    inputs = data["image"].to(device)
    target_availabilities = data["target_availabilities"].to(device)
    targets = data["target_positions"].to(device)
    # Forward pass
    preds, confidences = model(inputs)
    # skip compute loss if we are doing prediction
    loss = criterion(targets, preds, confidences, target_availabilities) if compute_loss else 0
    return loss, preds, confidences<jupyter_output><empty_output><jupyter_text>Init Model<jupyter_code># ==== INIT MODEL=================
model_name = cfg["model_params"]["model_name"]
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
#device = torch.device("cpu")
model = LyftMultiModel(cfg)

#load weight if there is a pretrained model
weight_path = f"{model_name}_final.pth"
if weight_path:
    model.load_state_dict(torch.load(weight_path))
model.to(device)
optimizer = optim.Adam(model.parameters(), lr=cfg["model_params"]["lr"])

# load weight if there is a pretrained model
opt_weight_path = f"{model_name}_optimizer_final.pth"
if opt_weight_path:
    optimizer.load_state_dict(torch.load(opt_weight_path))

print(f'device {device}')<jupyter_output>device cuda:0<jupyter_text>Training<jupyter_code>%%time
if cfg["model_params"]["train"]:
    tr_it = iter(train_dataloader)
    n_steps = cfg["train_params"]["steps"]
    progress_bar = tqdm.notebook.tqdm(range(1, 1 + n_steps), mininterval=5.)
    losses = []
    iterations = []
    metrics = []
    memorys = []
    times = []
    model_name = cfg["model_params"]["model_name"]
    update_steps = cfg['train_params']['update_steps']
    checkpoint_steps = cfg['train_params']['checkpoint_steps']
    t_start = time.time()
    torch.set_grad_enabled(True)
        
    for i in progress_bar:
        try:
            data = next(tr_it)
        except StopIteration:
            tr_it = iter(train_dataloader)
            data = next(tr_it)
        model.train()  

        optimizer.zero_grad()

        if cfg["train_params"]["precision"] == True:
            scaler = torch.cuda.amp.GradScaler()
            with torch.cuda.amp.autocast():
                loss, _, _ = forward(data, model, device)

            # Scales the loss, and calls backward()
            # to create scaled gradients
            scaler.scale(loss).backward()

            # Unscales gradients and calls
            # or skips optimizer.step()
            scaler.step(optimizer)

            # Updates the scale for next iteration
            scaler.update()
        else:
            
            loss, _, _ = forward(data, model, device)
            loss.backward()
            optimizer.step()
        
        loss_v = loss.item()
        losses.append(loss_v)
        
        if i % update_steps == 0:
            mean_losses = np.mean(losses)
            timespent = (time.time() - t_start) / 60
            print('i: %5d'%i,
                  'loss: %10.5f'%loss_v, 'loss(avg): %10.5f'%mean_losses, 
                  '%.2fmins'%timespent, end=' | ')
            mem = memory()
            if i % checkpoint_steps == 0:
                torch.save(model.state_dict(), f'{model_name}_{i}.pth')
                torch.save(optimizer.state_dict(), f'{model_name}_optimizer_{i}.pth')
            iterations.append(i)
            metrics.append(mean_losses)
            memorys.append(mem)
            times.append(timespent)

    torch.save(model.state_dict(), f'{model_name}_final.pth')
    torch.save(optimizer.state_dict(), f'{model_name}_optimizer_final.pth')
    results = pd.DataFrame({
        'iterations': iterations, 
        'metrics (avg)': metrics,
        'elapsed_time (mins)': times,
        'memory (GB)': memorys,
    })
    results.to_csv(f'train_metrics_{model_name}_{n_steps}_{runtime}.csv', index=False)
    print(f'Total training time is {(time.time() - t_start) / 60} mins')
    memory()
    display(results)<jupyter_output><empty_output><jupyter_text>Loss Curve<jupyter_code>plt.plot(np.arange(len(losses)), losses, label="train loss")
plt.legend()
plt.show()
%%time
if cfg["model_params"]["predict"]:
    
    model.eval()
    torch.set_grad_enabled(False)

    # store information for evaluation
    future_coords_offsets_pd = []
    timestamps = []
    confidences_list = []
    agent_ids = []
    memorys_pred = []
    t0 = time.time()
    times_pred = []
    iterations_pred = []

    for i, data in enumerate(tqdm.notebook.tqdm(test_dataloader, mininterval=5.)):
        
        _, preds, confidences = forward(data, model, device, compute_loss=False)
    
        preds = torch.einsum('bmti,bji->bmtj', 
                             preds.double(), 
                             data["world_from_agent"].to(device)[:, :2, :2]).cpu().numpy()
    
        future_coords_offsets_pd.append(preds.copy())
        confidences_list.append(confidences.cpu().numpy().copy())
        timestamps.append(data["timestamp"].numpy().copy())
        agent_ids.append(data["track_id"].numpy().copy()) 
        
        if i%50 == 0:
            t = ((time.time() - t0) / 60)
            print('%4d'%i, '%6.2fmins'%t, end=' | ')
            mem = memory()
            iterations_pred.append(i)
            memorys_pred.append(mem)
            times_pred.append(t)

    print('Total timespent: %6.2fmins'%((time.time() - t0) / 60))
    memory()
%%time
# create submission to submit to Kaggle
pred_path = 'submission.csv'
write_pred_csv(
    pred_path,
    timestamps=np.concatenate(timestamps),
    track_ids=np.concatenate(agent_ids),
    coords=np.concatenate(future_coords_offsets_pd),
    confs=np.concatenate(confidences_list),
)
df_sub = pd.read_csv(pred_path)
display(df_sub)<jupyter_output><empty_output>
StarcoderdataJupyterStructuredCleanDedup