train_path = './data/formatted_trainval/shanghaitech_part_A_patches_9/train' train_gt_path = './data/formatted_trainval/shanghaitech_part_A_patches_9/train_den' val_path = './data/formatted_trainval/shanghaitech_part_A_patches_9/val' val_gt_path = './data/formatted_trainval/shanghaitech_part_A_patches_9/val_den' #training configuration num_epochs=2000 lr = 0.00001 momentum = 0.9 # load net net = MCNN() weights_normal_init(net, dev=0.01) net.to(device) net.train() optimizer = torch.optim.Adam(net.parameters(), lr=lr, weight_decay=1e-5) criterion = nn.MSELoss() if not os.path.exists(checkpoint_dir): os.mkdir(checkpoint_dir) print('Loading training and validation datasets') #train_data_loader = ImageDataLoader(train_path, train_gt_path, shuffle=True, gt_downsample=True, pre_load=True) #val_data_loader = ImageDataLoader(val_path, val_gt_path, shuffle=False, gt_downsample=True, pre_load=True) train_dataset = Shanghai_Dataset(train_path, train_gt_path, gt_downsample=True, pre_load=True) train_data_loader = DataLoader(train_dataset, batch_size=1, shuffle=True, num_workers=0) val_dataset = Shanghai_Dataset(val_path, val_gt_path, gt_downsample=True, pre_load=True) val_data_loader = DataLoader(train_dataset, batch_size=1, shuffle=False, num_workers=0) best_mae = sys.maxint
pin_memory=True) else: train_loader = torch.utils.data.DataLoader(train_set, batch_size=batch_size, shuffle=True) # Initialize model, loss function, and optimizer model = MCNN() criterion = torch.nn.MSELoss() optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) # Iterate through training set and return input image and ground truth density map for batch_idx, (img, gt_dmap) in enumerate(train_loader): # Make sure we are in training mode model.train() # Torch accumulates gradients so we must zero these out for each batch optimizer.zero_grad() # Place data and labels in variable to track gradients and place on the GPU if available img = Variable(img.float()) gt_dmap = Variable(gt_dmap.float()) if torch.cuda.is_available(): img.cuda() pred_dmap.cuda() pred_map = model(img) import pdb pdb.set_trace()