def train_batch(batch_size=100): # Cuda or CPU: device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # Set params: params = Params() params.num_epoch = 35 params.lr = 1e-3 # Load data: X = load_data_batch_v2().to(device) print('Training data:', X.shape) # Init model: model = create_model(params) model.to(device) print(model) #Optimizer params: optimizer_params = [] for key, module in dict(model.named_parameters()).items(): if key == 'reproject.weight': optimizer_params.append({ 'params': module, 'weight_decay': 0.0, 'lr': 0.01 }) else: optimizer_params.append({'params': module}) # Loss and optimizer: criterion = torch.nn.MSELoss() # optimizer = torch.optim.Adam(model.parameters(), lr=params.lr, weight_decay=params.weight_decay) optimizer = torch.optim.Adam(optimizer_params, lr=params.lr, weight_decay=params.weight_decay) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=15, gamma=0.1) # Training model.train() n = X.shape[0] for ei in range(params.num_epoch): start = time.time() for i in range(0, n, batch_size): count = min(batch_size, n - i) X_batch = X[i:i + count] # Forward pass: optimizer.zero_grad() # X_pred, joints3d_pred = model.forward(X_batch) X_pred = model.forward(X_batch) # Compute the loss: loss = criterion(X_pred, X_batch) # Backward pass: loss.backward() # Update params: optimizer.step() # Log: elapsed_time = time.time() - start lr_log = ["%.5f" % g['lr'] for g in optimizer.param_groups] print("[EPOCH]: %i, [LOSS]: %.6f, [elapsed]: %.4fs, [lr]: %s" % (ei, loss.item(), elapsed_time, lr_log)) # Scheduler: scheduler.step() # Save model: model_path = '/home/darkalert/KazendiJob/DLab/multi-view-pose-3d/models/v2-big-norm-weighted-mse' + str( loss.item()) + '.pth' torch.save(model.state_dict(), model_path) print('Model has been saved to ', model_path) # Test: joint3d = model.predict(X[0]).view(17, 3).cpu().numpy() print('Xi:', X[100]) print(joint3d) draw_skeleton_3d(joint3d)
def train(): # Cuda or CPU: device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # Set params: params = Params() params.num_epoch = 1 params.lr = 1e-4 # Load data: X = load_data_batch_v1().to(device) print('Training data:', X.shape) # Init model: model = create_model(params) model.to(device) print(model) # # Optimizer params: # optimizer_params = [] # for key,module in dict(model.named_parameters()).items(): # if key == 'reproject.weight': # optimizer_params.append({'params': module,'weight_decay':0.0,'lr':params.lr}) # else: # optimizer_params.append({'params': module}) # Training: criterion = torch.nn.MSELoss() optimizer = torch.optim.Adam(model.parameters(), lr=params.lr, weight_decay=params.weight_decay) # optimizer = torch.optim.Adam(optimizer_params, lr=params.lr, weight_decay=params.weight_decay) model.train() for ei in range(params.num_epoch): for i in range(X.shape[0]): # Forward pass: optimizer.zero_grad() X_pred = model.forward(X[i]) # Compute the loss: loss = criterion(X_pred, X[i]) if i % 10 == 0: print("[ITER]: %i, [LOSS]: %.6f" % (i, loss.item())) # Backward pass: loss.backward() # for i,p in enumerate(model.parameters()): # print (i,p.grad.shape) # Update params: optimizer.step() print("[EPOCH]: %i, [LOSS]: %.6f" % (ei, loss.item())) # Save model: model_path = '/home/darkalert/KazendiJob/DLab/multi-view-pose-3d/models/last_model_v2_' + str( loss.item()) + '.pth' torch.save(model.state_dict(), model_path) print('Model has been saved to ', model_path) # Test: joint3d = model.predict(X[0]).view(17, 3).cpu().numpy() print('Xi:', X[i]) print(joint3d) draw_skeleton_3d(joint3d)