import os,time import torch,wandb from tqdm import trange from einops import rearrange import vis,geometry from copy import deepcopy import numpy as np import piqa,kornia from torchvision.utils import make_grid from einops import rearrange, repeat from our_models import ch_sec from torch.cuda.amp import autocast, GradScaler import torch.nn.functional as F import torch.nn as nn import getpass from glob import glob import data import our_models as models def to_gpu(ob): return {k: to_gpu(torch.tensor(v)) for k, v in ob.items()} if isinstance(ob, dict) else ob.cuda(non_blocking=True) def train(run,until_img=25,until_vid=100,until_save=500,optim=None,single_data=None): def loss_fn(model_out, gt, data_sample,model,step): rays = lambda x,y: torch.stack([x[i,:,y[i].long()] for i in range(len(x))]) losses = { } losses["metrics/cam_pointmap_loss"] = (data_sample["points_cam"] - model_out["points_cam"]).square().mean()*3e2 losses["metrics/link_pointmap_loss"] = (data_sample["points_link"] - model_out["points_link"]).square().mean()*4e2 losses["metrics/segs"] = (data_sample["segs"].float() - model_out["segs"]).square().mean()*1e2 print(losses) return losses losses_agg=[] optim = torch.optim.Adam(lr=run.args.lr, params=run.model.parameters()) get_loader = lambda dataset: torch.utils.data.DataLoader(dataset, batch_size=run.args.batch_size, num_workers=min(run.args.n_workers,run.args.batch_size),pin_memory=True) datasets=[data.PyBulletFolder("/data/cameron/robot_calibration_testing/fulldiversecam_redo/",val=val) for val in range(2)] dataloaders=[iter(get_loader(x)) for x in datasets] # Train loop step=0 for step_ in trange(run.args.n_train_steps, desc="Fitting"): # train until user interruption of=0 val=step%101==0 and step and not of suffix="val/" if val else "train/" # Get data if of:print("overfitting on sample") if step_==0 or not of: if not run.args.online:data_sample= to_gpu(next(dataloaders[val])) else: try: data_sample= to_gpu(next(dataloaders[val])) except: print("skipping bad load") step-=1 continue # Run model and calculate losses total_loss = 0. with torch.set_grad_enabled(not val):out=run.model(data_sample) # Mask background for pointmaps in input and pred also add background channel for k in ["points_cam","points_link"]: for d in [data_sample,out]: d[k] = d[k]*data_sample["fg"] losses = loss_fn(out, data_sample, data_sample,run.model,step) for loss_name, loss in losses.items(): wandb.log({loss_name+suffix: loss.item()}, step=step) total_loss += loss wandb.log({"loss": total_loss.item()}, step=step) if not val :total_loss.backward();optim.step(); optim.zero_grad(); with torch.no_grad(): wandb_imgs=None if step%until_img==0 or val or 0: wandb_imgs=vis.wandb_summary( 0, out, data_sample, data_sample, None,step=step,prefix=suffix) if step%until_save == 0 and step and run.args.save_model: # save model print(f"Saving to {run.save_dir}"); torch.save({ 'step': step, 'model_state_dict': run.model.state_dict(), }, os.path.join(run.save_dir, f"checkpoint.pt")) step+=1 # Data/args setup and run import argparse parser = argparse.ArgumentParser(description='simple training job') # logging parameters parser.add_argument('-n','--name', type=str,default="",required=False,help="wandb training name") parser.add_argument('-c','--init_ckpt', type=str,default=None,required=False,help="File for checkpoint loading. If folder specific, will use latest .pt file") parser.add_argument('-o','--online', default=False, action='store_true') parser.add_argument('-s','--save_model', default=True, action='store_true') parser.add_argument('--viser', default=False, action='store_true') parser.add_argument('--save_opt_vis', default=False, action='store_true') # data/training parameters parser.add_argument('-d','--dataset', type=str,default="hydrant") parser.add_argument('--imgpath', type=str,default="") parser.add_argument('--dirpath', type=str,default="/data/cameron/robot_calibration_testing/justjoints/")#rotation_gripper_dataset/") parser.add_argument('-b','--batch_size', type=int,default=1,help="number of videos/sequences per training step") parser.add_argument('-v','--vid_len', type=int,default=6,help="video length or number of images per batch") parser.add_argument('--n_workers',type=int,default=8,help="number of workers per dataloader") parser.add_argument('--until_save',type=int,default=500,help="number of steps until model save") parser.add_argument('--lr',type=float,default=1e-4,help="learning rate") parser.add_argument('--n_train_steps',type=int,default=10000,help="n train steps ") parser.add_argument('--overfit', default=True, action='store_true',help="Whether to overfit on a single scene") parser.add_argument('--until_img', type=int,default=100,help="Number of steps until image summary. ") parser.add_argument('--sf', type=float,default=1,help="Image resolution scale factor (fractional is cheaper)") parser.add_argument('--load_save', default=False, action='store_true',help="Whether to load the previously saved data if overfitting (to avoid running flow again)") # model parameters parser.add_argument('--noncanon', default=False, action='store_true',help="use canon frame model or naive baseline") import pybullet as p import pybullet_data def make_run(args=None,val=False): args = parser.parse_args(args) self = argparse.Namespace() user = getpass.getuser() print(f"user={user}") # Wandb init run = wandb.init(entity="cameronsmithbusiness",project="biasing",mode="online" if args.online else "disabled",name=args.name,dir=f"/tmp/wandb") wandb.run.log_code(".") self.save_dir = "/tmp/"+args.name#os.path.join(os.environ.get('LOGDIR', "") , run.name) os.makedirs(self.save_dir,exist_ok=True) wandb.save(os.path.join(self.save_dir, "checkpoint*")) wandb.save(os.path.join(self.save_dir, "video*")) self.args=args self.wandb=run if args.viser: self.viser_server=viser.ViserServer() # Pybullet setup physicsClient = p.connect(p.DIRECT);p.setGravity(0, 0, -9.8) p.setAdditionalSearchPath(pybullet_data.getDataPath()) self.args.rob_id = p.loadURDF("franka_panda/panda.urdf") # Make model and load checkpoint self.model = models.CanonRobotTrajPredTransfEuclidean(args).cuda() if args.init_ckpt is not None: ckpt_file = args.init_ckpt if os.path.isfile(os.path.expanduser(args.init_ckpt)) else max(glob(os.path.join(args.init_ckpt,"*.pt")), key=os.path.getctime) self.model.load_state_dict(torch.load(ckpt_file)["model_state_dict"],strict=False) return self if __name__=="__main__": run = make_run() torch.autograd.set_detect_anomaly(False) train(run,until_save=run.args.until_save, until_vid=100 if not run.args.overfit else 300, until_img=run.args.until_img)