def main(args): model = resnet.ResNet(torchvision.models.resnet.Bottleneck, [3, 4, 6, 3], args.num_classes) saved_state_dict = torch.load(args.saved_model) transformations = transforms.Compose([ transforms.Resize((args.image_size, args.image_size)), transforms.ToTensor() ]) if args.gpu[0] >= 0: cudnn.enabled = True softmax = nn.Softmax().cuda() model.cuda() else: softmax = nn.Softmax() load_filtered_state_dict(model, saved_state_dict, ignore_layer=[], reverse=True) test_x, test_y, classes_names = get_dataset(args.test_data_dir) test_dataset = DataWrapper(test_x, test_y, transformations, augumentation=False) test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=args.batch_size, shuffle=False, num_workers=1) classes, indices = np.unique(test_y, return_index=True) #n = (test_dataset.__len__() + args.batch_size - 1) / args.batch_size * args.batch_size n = test_dataset.__len__() y_pred = np.zeros((n)) y = np.zeros((n)) count = 0 for i, (images, labels, names) in enumerate(test_loader): images = Variable(images) labels = Variable(labels) if args.gpu[0] >= 0: images = images.cuda() labels = labels.cuda() label_pred = model(images) label_pred = softmax(label_pred) n = images.size()[0] _, label_pred = label_pred.topk(1, 1, True, True) y_pred[count:count + n] = label_pred.view(-1).cpu().numpy() y[count:count + n] = labels.data.cpu().numpy() count += n plot(y, y_pred, classes_names)
def main(args): cudnn.enabled = True print('Loading data.') transformations = transforms.Compose([ transforms.Resize(240), transforms.RandomCrop(224), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) train_x, train_y, classes_names = get_dataset(args.trainning_data_dir) test_x, test_y, _ = get_dataset(args.validation_data_dir) num_classes = len(classes_names) trainning_dataset = DataWrapper(train_x, train_y, transformations) eval_dataset = DataWrapper(test_x, test_y, transformations) train_loader = torch.utils.data.DataLoader(dataset=trainning_dataset, batch_size=args.batch_size, shuffle=True, num_workers=16) eval_loader = torch.utils.data.DataLoader(dataset=eval_dataset, batch_size=args.batch_size, shuffle=True, num_workers=16) n = trainning_dataset.__len__() print(n) # ResNet50 structure model = resnet.ResNet(torchvision.models.resnet.Bottleneck, [3, 4, 6, 3], num_classes) if args.saved_model: print('Loading model.') saved_state_dict = torch.load(args.saved_model) # 'origin model from pytorch' if 'resnet' in args.saved_model: load_filtered_state_dict(model, saved_state_dict, ignore_layer=[], reverse=False) else: load_filtered_state_dict(model, saved_state_dict, ignore_layer=[]) crossEntropyLoss = nn.CrossEntropyLoss().cuda() softmax = nn.Softmax().cuda() optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) # multi-gpu model = nn.DataParallel(model, device_ids=[0, 1]) model.cuda() Save_model = SaveBestModel(save_dir=args.save_path) Writer = SummaryWriter() step = 0 for epoch in range(args.num_epochs): evaluate(eval_loader, model, Writer, step, Save_model, epoch) step = train(train_loader, model, crossEntropyLoss, optimizer, Writer, args.batch_size, epoch, step, n)