simclr_ds_training.py

for epoch in range(20):
    
    stime = time.time()
    print("=============== Epoch : %3d ==============="%(epoch+1))
    
    loss_sublist = np.array([])
    acc_sublist = np.array([])
    
    #iter_num = 0
    dsmodel.train()
    
    dsoptimizer.zero_grad()
    
    for x,y in dl:
        x = x.squeeze().to(device = 'cuda:0', dtype = torch.float)
        y = y.to(device = 'cuda:0')
        
        z = dsmodel(x)
        
        dsoptimizer.zero_grad()
        
        tr_loss = loss_fn(z,y)
        tr_loss.backward()

        preds = torch.exp(z.cpu().data)/torch.sum(torch.exp(z.cpu().data))
        
        dsoptimizer.step()
        
        loss_sublist = np.append(loss_sublist, tr_loss.cpu().data)
        acc_sublist = np.append(acc_sublist,np.array(np.argmax(preds,axis=1)==y.cpu().data.view(-1)).astype('int'),axis=0)
        
    print('ESTIMATING TRAINING METRICS.............')
    
    print('TRAINING BINARY CROSSENTROPY LOSS: ',np.mean(loss_sublist))
    print('TRAINING BINARY ACCURACY: ',np.mean(acc_sublist))
    
    tr_ep_loss.append(np.mean(loss_sublist))
    tr_ep_acc.append(np.mean(acc_sublist))
    
    print('ESTIMATING VALIDATION METRICS.............')
    
    dsmodel.eval()
    
    loss_sublist = np.array([])
    acc_sublist = np.array([])
    
    with torch.no_grad():
        for x,y in vdl:
            x = x.squeeze().to(device = 'cuda:0', dtype = torch.float)
            y = y.to(device = 'cuda:0')
            z = dsmodel(x)

            val_loss = loss_fn(z,y)

            preds = torch.exp(z.cpu().data)/torch.sum(torch.exp(z.cpu().data))

            loss_sublist = np.append(loss_sublist, val_loss.cpu().data)
            acc_sublist = np.append(acc_sublist,np.array(np.argmax(preds,axis=1)==y.cpu().data.view(-1)).astype('int'),axis=0)
    
    print('VALIDATION BINARY CROSSENTROPY LOSS: ',np.mean(loss_sublist))
    print('VALIDATION BINARY ACCURACY: ',np.mean(acc_sublist))
    
    val_ep_loss.append(np.mean(loss_sublist))
    val_ep_acc.append(np.mean(acc_sublist))
    
    lr_scheduler.step()
    
    dg.on_epoch_end()
    
    if np.mean(loss_sublist) <= min_val_loss:
        min_val_loss = np.mean(loss_sublist) 
        print('Saving model...')
        torch.save({'model_state_dict': dsmodel.state_dict(),
                'optimizer_state_dict': dsoptimizer.state_dict()}, 
               '/content/saved_models/cifar10_rn50_p128_sgd0p01_decay0p98_all_lincls.pt')
    
    print("Time Taken : %.2f minutes"%((time.time()-stime)/60.0))