def main(): # Prepare train dataset and dataloader train_ds = PascalVOCDataset('./data', 'TRAIN', keep_difficult=keep_difficult) train_loader = torch.utils.data.DataLoader(train_ds, batch_size=batch_size, shuffle=True, collate_fn=train_ds.collate_fn, # note that we're passing the collate function here num_workers=num_workers, pin_memory=True) n_classes = len(train_ds.label_map()) start_epoch = 0 # Initialize model model = SSD300(n_classes=n_classes) # Load checkpoint if existed checkpoint = None if checkpoint_path is not None and os.path.exists(checkpoint_path): checkpoint = torch.load(checkpoint_path) start_epoch = checkpoint['epoch'] + 1 print('Load checkpoint from epoch %d.\n' % checkpoint['epoch']) if checkpoint is not None: model.load_state_dict(checkpoint['model_state_dict']) model.to(device) model.train() # Initialize the optimizer, with twice the default learning rate for biases, as in the original Caffe repo biases = list() not_biases = list() for param_name, param in model.named_parameters(): if param.requires_grad: if param_name.endswith('.bias'): biases.append(param) else: not_biases.append(param) optimizer = torch.optim.SGD(params=[{'params': biases, 'lr': 2 * lr}, {'params': not_biases}], lr=lr, momentum=momentum, weight_decay=weight_decay) if checkpoint is not None: optimizer.load_state_dict(checkpoint['optimizer_state_dict']) criterion = MultiBoxLoss(priors_cxcy=model.priors_cxcy).to(device) # Calculate total number of epochs to train and the epochs to decay learning rate at (i.e. convert iterations to epochs) epochs = iterations // (len(train_ds) // batch_size) decay_lr_at_epochs = [it // (len(train_ds) // batch_size) for it in decay_lr_at] # Epochs for epoch in range(start_epoch, epochs): # Decay learning rate at particular epochs if epoch in decay_lr_at_epochs: utils.adjust_learning_rate(optimizer, decay_lr_to) # One epoch's training train(train_loader, model, criterion, optimizer, epoch) # Save checkpoint utils.save_checkpoint(checkpoint_path, model, optimizer, epoch)
from tqdm import tqdm device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # Parameters keep_difficult = True # difficult ground truth objects must always be considered in mAP calculation, because these objects DO exist! batch_size = 64 workers = 4 checkpoint_path = 'ssd300.pt' # Load test data test_dataset = PascalVOCDataset('./data', split='TEST', keep_difficult=keep_difficult) test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False, collate_fn=test_dataset.collate_fn, num_workers=workers, pin_memory=True) n_classes = len(test_dataset.label_map()) # Load model checkpoint that is to be evaluated checkpoint = torch.load(checkpoint_path) model = SSD300(n_classes) model.load_state_dict(checkpoint['model_state_dict']) model = model.to(device) # Switch to eval mode model.eval() def evaluate(test_loader, model): """ Evaluate.