def main(): # Model Initialize m = createModel().cuda() if opt.loadModel: print('Loading Model from {}'.format(opt.loadModel)) m.load_state_dict(torch.load(opt.loadModel)) if not os.path.exists("../exp/{}/{}".format(opt.dataset, opt.expID)): try: os.mkdir("../exp/{}/{}".format(opt.dataset, opt.expID)) except FileNotFoundError: os.mkdir("../exp/{}".format(opt.dataset)) os.mkdir("../exp/{}/{}".format(opt.dataset, opt.expID)) else: print('Create new model') if not os.path.exists("../exp/{}/{}".format(opt.dataset, opt.expID)): try: os.mkdir("../exp/{}/{}".format(opt.dataset, opt.expID)) except FileNotFoundError: os.mkdir("../exp/{}".format(opt.dataset)) os.mkdir("../exp/{}/{}".format(opt.dataset, opt.expID)) criterion = torch.nn.MSELoss().cuda() if opt.optMethod == 'rmsprop': optimizer = torch.optim.RMSprop(m.parameters(), lr=opt.LR, momentum=opt.momentum, weight_decay=opt.weightDecay) elif opt.optMethod == 'adam': optimizer = torch.optim.Adam(m.parameters(), lr=opt.LR) else: raise Exception writer = SummaryWriter('.tensorboard/{}/{}'.format(opt.dataset, opt.expID)) # Prepare Dataset if opt.dataset == 'coco': train_dataset = coco.Mscoco(train=True) val_dataset = coco.Mscoco(train=False) train_dataset = coco.Mscoco(train=True) val_dataset = coco.Mscoco(train=False) train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=opt.trainBatch, shuffle=True, num_workers=opt.nThreads, pin_memory=True) val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=opt.validBatch, shuffle=False, num_workers=opt.nThreads, pin_memory=True) # Model Transfer m = torch.nn.DataParallel(m).cuda() # Start Training for i in range(opt.nEpochs): opt.epoch = i print('############# Starting Epoch {} #############'.format( opt.epoch)) loss, acc = train(train_loader, m, criterion, optimizer, writer) print('Train-{idx:d} epoch | loss:{loss:.8f} | acc:{acc:.4f}'.format( idx=opt.epoch, loss=loss, acc=acc)) opt.acc = acc opt.loss = loss m_dev = m.module if i % opt.snapshot == 0: torch.save( m_dev.state_dict(), '../exp/{}/{}/model_{}.pkl'.format(opt.dataset, opt.expID, opt.epoch)) torch.save( opt, '../exp/{}/{}/option.pkl'.format(opt.dataset, opt.expID, opt.epoch)) torch.save( optimizer, '../exp/{}/{}/optimizer.pkl'.format(opt.dataset, opt.expID)) loss, acc = valid(val_loader, m, criterion, optimizer, writer) print('Valid-{idx:d} epoch | loss:{loss:.8f} | acc:{acc:.4f}'.format( idx=i, loss=loss, acc=acc)) ''' if opt.dataset != 'mpii': with torch.no_grad(): mAP, mAP5 = prediction(m) print('Prediction-{idx:d} epoch | mAP:{mAP:.3f} | mAP0.5:{mAP5:.3f}'.format( idx=i, mAP=mAP, mAP5=mAP5 )) ''' writer.close()
def main(): print(opt) # Model Initialize m = createModel().cuda() if opt.loadModel: print('Loading Model from {}'.format(opt.loadModel)) ''' ckp = torch.load(opt.loadModel) for name,param in m.state_dict().items(): if name in ckp: ckp_param = ckp[name] if ckp_param.shape == param.shape: param.copy_(ckp_param) print(name, 'copy successfully') else: print(name, 'shape is inconsistent with checkpoint') else: print(name, 'can not be found in checkpoint') ''' m.load_state_dict(torch.load(opt.loadModel)) if not os.path.exists("../exp/{}/{}".format(opt.dataset, opt.expID)): try: os.mkdir("../exp/{}/{}".format(opt.dataset, opt.expID)) except FileNotFoundError: os.mkdir("../exp/{}".format(opt.dataset)) os.mkdir("../exp/{}/{}".format(opt.dataset, opt.expID)) else: print('Create new model') # import pdb;pdb.set_trace() if not os.path.exists("../exp/{}/{}".format(opt.dataset, opt.expID)): try: os.mkdir("../exp/{}/{}".format(opt.dataset, opt.expID)) except FileNotFoundError: os.mkdir("../exp/{}".format(opt.dataset)) os.mkdir("../exp/{}/{}".format(opt.dataset, opt.expID)) criterion = torch.nn.MSELoss().cuda() if opt.optMethod == 'rmsprop': optimizer = torch.optim.RMSprop(m.parameters(), lr=opt.LR, momentum=opt.momentum, weight_decay=opt.weightDecay) elif opt.optMethod == 'adam': optimizer = torch.optim.Adam(m.parameters(), lr=opt.LR) elif opt.optMethod == 'rmsprop_refine': print('opt rmsprop refine') optimizer = torch.optim.RMSprop([{ "params": m.preact.parameters(), "lr": 1e-5 }, { "params": m.suffle1.parameters(), "lr": 1e-5 }, { "params": m.duc1.parameters(), "lr": 1e-5 }, { "params": m.duc2.parameters(), "lr": 1e-5 }, { "params": m.conv_out.parameters(), "lr": 1e-4 }], lr=opt.LR, momentum=opt.momentum, weight_decay=opt.weightDecay) else: raise Exception if opt.loadOptimizer: optimizer = torch.load(opt.loadOptimizer) writer = SummaryWriter('.tensorboard/{}/{}'.format(opt.dataset, opt.expID)) # Prepare Dataset if opt.dataset == 'coco': train_dataset = coco.Mscoco(train=True) val_dataset = coco.Mscoco(train=False) elif opt.dataset == 'h36m': train_dataset = h36m.H36M(train=True) val_dataset = h36m.H36M(train=False) ''' train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=opt.trainBatch, shuffle=False, sampler=RandomSampler(train_dataset, replacement=True, num_samples=len(train_dataset)//10), num_workers=opt.nThreads, pin_memory=True) val_loader = torch.utils.data.DataLoader( val_dataset, batch_size=opt.validBatch, shuffle=False, sampler=RandomSampler(val_dataset, replacement=True, num_samples=len(val_dataset)//10), num_workers=opt.nThreads, pin_memory=True) ''' train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=opt.trainBatch, shuffle=True, num_workers=opt.nThreads, pin_memory=True) val_loader = torch.utils.data.DataLoader( val_dataset, batch_size=opt.validBatch, shuffle=False, sampler=RandomSampler(val_dataset, replacement=True, num_samples=len(val_dataset) // 8), num_workers=opt.nThreads, pin_memory=True) #show_loader_image('train_check_images', train_loader, joint_names=train_dataset.joint_names) #show_loader_image('valid_check_images', val_loader, joint_names=val_dataset.joint_names) #return # Model Transfer m = torch.nn.DataParallel(m).cuda() # Start Training for i in range(opt.nEpochs): opt.epoch = i print('############# Starting Epoch {} #############'.format( opt.epoch)) loss, acc = train(train_loader, m, criterion, optimizer, writer) print('Train-{idx:d} epoch | loss:{loss:.8f} | acc:{acc:.4f}'.format( idx=opt.epoch, loss=loss, acc=acc)) opt.acc = acc opt.loss = loss m_dev = m.module if i % opt.snapshot == 0: torch.save( m_dev.state_dict(), '../exp/{}/{}/model_{}.pkl'.format(opt.dataset, opt.expID, opt.epoch)) torch.save( opt, '../exp/{}/{}/option.pkl'.format(opt.dataset, opt.expID, opt.epoch)) torch.save( optimizer, '../exp/{}/{}/optimizer.pkl'.format(opt.dataset, opt.expID)) loss, acc = valid(val_loader, m, criterion, optimizer, writer) print('Valid-{idx:d} epoch | loss:{loss:.8f} | acc:{acc:.4f}'.format( idx=i, loss=loss, acc=acc)) ''' if opt.dataset != 'mpii': with torch.no_grad(): mAP, mAP5 = prediction(m) print('Prediction-{idx:d} epoch | mAP:{mAP:.3f} | mAP0.5:{mAP5:.3f}'.format( idx=i, mAP=mAP, mAP5=mAP5 )) ''' writer.close()
def main(): # Model Initialize m = createModel().cuda() m.load_state_dict(torch.load('duc_se.pth')) criterion = torch.nn.MSELoss().cuda() if opt.optMethod == 'rmsprop': optimizer = torch.optim.RMSprop(m.parameters(), lr=opt.LR, momentum=opt.momentum, weight_decay=opt.weightDecay) elif opt.optMethod == 'adam': optimizer = torch.optim.Adam(m.parameters(), lr=opt.LR) else: raise Exception writer = SummaryWriter('.tensorboard/{}/{}'.format(opt.dataset, opt.expID)) # Prepare Dataset if opt.dataset == 'coco': train_dataset = coco.Mscoco(train=True) val_dataset = coco.Mscoco(train=False) train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=opt.trainBatch, shuffle=True, num_workers=opt.nThreads, pin_memory=True) val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=opt.validBatch, shuffle=False, num_workers=opt.nThreads, pin_memory=True) # Model Transfer m = torch.nn.DataParallel(m).cuda() # Start Training for i in range(opt.nEpochs): opt.epoch = i print('############# Starting Epoch {} #############'.format( opt.epoch)) loss, acc = train(train_loader, m, criterion, optimizer, writer) print('Train-{idx:d} epoch | loss:{loss:.8f} | acc:{acc:.4f}'.format( idx=opt.epoch, loss=loss, acc=acc)) opt.acc = acc opt.loss = loss m_dev = m.module if i % opt.snapshot == 0: torch.save( m_dev.state_dict(), '../exp/{}/{}/model_{}.pkl'.format(opt.dataset, opt.expID, opt.epoch)) torch.save( opt, '../exp/{}/{}/option.pkl'.format(opt.dataset, opt.expID, opt.epoch)) torch.save( optimizer, '../exp/{}/{}/optimizer.pkl'.format(opt.dataset, opt.expID)) loss, acc = valid(val_loader, m, criterion, optimizer, writer) print('Valid-{idx:d} epoch | loss:{loss:.8f} | acc:{acc:.4f}'.format( idx=i, loss=loss, acc=acc)) writer.close()
def main(): # 初始化模型 m = createModel().cuda() # 加载已经训练道模型接着训练 if opt.loadModel: print("加载已训练道模型".format(opt.loadModel)) m.load_state_dict(torch.load(opt.loadModel)) pass # 新建模型进行训练 else: print('新建一个模型') if not os.path.exists("../exp/{}/{}".format(opt.dataset, opt.expID)): try: os.mkdir("../exp/{}/{}".format(opt.dataset, opt.expID)) except FileNotFoundError: os.mkdir("../exp/{}".format(opt.dataset)) os.mkdir("../exp/{}/{}".format(opt.dataset, opt.expID)) criterion = torch.nn.MSELoss().cuda() # 定义优化器 if opt.optMethod == 'adam': optimizer = torch.optim.Adam( m.parameters(), lr=opt.LR ) elif opt.optMethod == 'rmsprop': optimizer = torch.optim.RMSprop(m.parameters(), lr=opt.LR, momentum=opt.momentum, weight_decay=opt.weightDecay) else: raise Exception writer = SummaryWriter( '.tensorboard/{}/{}'.format(opt.dataset, opt.expID) ) # 准备训练数据 if opt.dataset == 'coco': train_dataset =coco.Mscoco(train=True) val_dataset = coco.Mscoco(train=False) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=opt.trainBatch, shuffle=True, num_workers=opt.nThreads, pin_memory=True ) val_loader = torch.utils.data.DataLoader( val_dataset, batch_size=opt.validBatch, shuffle=False, num_workers=opt.nThreads, pin_memory=True ) m = torch.nn.DataParallel(m).cuda() # 开始训练 for i in range(opt.nEpochs): opt.epoch = i print('############## String Epoch ##########'.format(opt.epoch)) loss, acc = train(train_loader, m, criterion, optimizer, writer) print('Train-{idx:d} epoch | loss:{loss:.8f} | acc:{acc:.4f}'.format( idx=opt.epoch, loss=loss, acc=acc )) opt.acc = acc opt.loss = loss m_dev = m.module if i % opt.snapshot == 0: torch.save( m_dev.state_dict(), '../exp/{}/{}/model_{}.pkl'.format(opt.dataset, opt.expID, opt.epoch) ) torch.save( opt, '../exp/{}/{}/option.pkl'.format(opt.dataset, opt.expID, opt.epoch) ) torch.save( optimizer, '../exp/{}/{}/optimizer.pkl'.format(opt.dataset, opt.expID) ) loss, acc = valid(val_loader, m, criterion, optimizer, writer) print( 'Valid-{idx:d} epoch | loss:{loss:.8f} | acc:{acc:.4f}'.format( idx=i, loss=loss, acc=acc) ) writer.close()
def main(): # Model Initialize m = createModel().cuda() if opt.loadModel: print('Loading Model from {}'.format(opt.loadModel)) current_model_weight = m.state_dict() weight_save = torch.load(opt.loadModel) weight_save_changed = {} for k in weight_save: if 'conv_out.weight' in k or 'conv_out.bias' in k: print(k, 'not used') continue weight_save_changed[k] = weight_save[k] current_model_weight.update(weight_save_changed) m.load_state_dict(current_model_weight) if not os.path.exists("../exp/{}/{}".format(opt.dataset, opt.expID)): try: os.mkdir("../exp/{}/{}".format(opt.dataset, opt.expID)) except FileNotFoundError: os.mkdir("../exp") os.mkdir("../exp/{}".format(opt.dataset)) os.mkdir("../exp/{}/{}".format(opt.dataset, opt.expID)) else: print('Create new model') if not os.path.exists("../exp/{}/{}".format(opt.dataset, opt.expID)): try: os.mkdir("../exp/{}/{}".format(opt.dataset, opt.expID)) except FileNotFoundError: os.mkdir("../exp") os.mkdir("../exp/{}".format(opt.dataset)) os.mkdir("../exp/{}/{}".format(opt.dataset, opt.expID)) criterion = torch.nn.MSELoss().cuda() if opt.optMethod == 'rmsprop': optimizer = torch.optim.RMSprop(m.parameters(), lr=opt.LR) elif opt.optMethod == 'adam': optimizer = torch.optim.Adam(m.parameters(), lr=opt.LR) else: raise Exception writer = SummaryWriter('.tensorboard/{}/{}'.format(opt.dataset, opt.expID)) # Prepare Dataset if opt.dataset == 'coco': train_dataset = coco.Mscoco(train=True, img_folder=opt.img_folder_train, annot_file=opt.annot_file_train, nJoints=opt.nClasses) val_dataset = coco.Mscoco(train=False, img_folder=opt.img_folder_val, annot_file=opt.annot_file_val, nJoints=opt.nClasses) train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=opt.trainBatch, shuffle=True, num_workers=opt.nThreads, pin_memory=True) val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=opt.validBatch, shuffle=False, num_workers=opt.nThreads, pin_memory=True) # Model Transfer m = DataParallel(m).cuda() # Start Training for i in range(opt.nEpochs + 1): opt.epoch = i print('############# Starting Epoch {} #############'.format( opt.epoch)) loss, acc = train(train_loader, m, criterion, optimizer, writer) print('Train-{idx:d} epoch | loss:{loss:.8f} | acc:{acc:.4f}'.format( idx=opt.epoch, loss=loss, acc=acc)) opt.acc = acc opt.loss = loss m_dev = m.module if (i % opt.snapshot == 0 and i != 0) or i == opt.nEpochs: torch.save( m_dev.state_dict(), '../exp/{}/{}/model_{}.pkl'.format(opt.dataset, opt.expID, opt.epoch)) torch.save( opt, '../exp/{}/{}/option.pkl'.format(opt.dataset, opt.expID, opt.epoch)) # torch.save( # optimizer, '../exp/{}/{}/optimizer.pkl'.format(opt.dataset, opt.expID)) loss, acc = valid(val_loader, m, criterion, optimizer, writer) print('Valid-{idx:d} epoch | loss:{loss:.8f} | acc:{acc:.4f}'.format( idx=i, loss=loss, acc=acc)) # # Prediction Test # if opt.map and acc > 0.8: # with torch.no_grad(): # mAP, mAP5 = prediction_silence(m) # logger.info('Prediction-{idx:d} epoch | mAP:{mAP:.3f} | mAP0.5:{mAP5:.3f}'.format( # idx=i, # mAP=mAP, # mAP5=mAP5 # )) ''' if opt.dataset != 'mpii': with torch.no_grad(): mAP, mAP5 = prediction(m) print('Prediction-{idx:d} epoch | mAP:{mAP:.3f} | mAP0.5:{mAP5:.3f}'.format( idx=i, mAP=mAP, mAP5=mAP5 )) ''' writer.close()