()
| 37 | help='number of data loading workers (default: 4)') |
| 38 | |
| 39 | def main(): |
| 40 | global args |
| 41 | args = parser.parse_args() |
| 42 | log_file_name = args.save_name + time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime(time.time())) + '.log' |
| 43 | global logger |
| 44 | logger = create_logger(os.path.join(args.exp, log_file_name)) |
| 45 | logger.info("============ Initialized logger ============") |
| 46 | logger.info("\n".join("%s: %s" % (k, str(v)) |
| 47 | for k, v in sorted(dict(vars(args)).items()))) |
| 48 | logger.info("The experiment will be stored in %s\n" % args.exp) |
| 49 | logger.info("") |
| 50 | |
| 51 | # fix random seeds |
| 52 | torch.manual_seed(args.seed) |
| 53 | torch.cuda.manual_seed_all(args.seed) |
| 54 | np.random.seed(args.seed) |
| 55 | best_prec1 = 0 |
| 56 | |
| 57 | # network defined |
| 58 | checkpoint = torch.load(args.model) |
| 59 | model = models.__dict__[checkpoint['arch']](out=args.nmb_cluster, linear_eval=True, extra_mlp=True) |
| 60 | |
| 61 | # freeze the features layers |
| 62 | for param in model.parameters(): |
| 63 | param.requires_grad = False |
| 64 | for param in model.linear.parameters(): |
| 65 | param.requires_grad = True |
| 66 | |
| 67 | # load model |
| 68 | model = torch.nn.DataParallel(model) |
| 69 | model.load_state_dict(checkpoint['state_dict'], strict=False) |
| 70 | model.cuda() |
| 71 | cudnn.benchmark = True |
| 72 | |
| 73 | # define loss function |
| 74 | criterion = nn.CrossEntropyLoss() |
| 75 | |
| 76 | # train & val dataloader |
| 77 | traindir = os.path.join(args.data, 'train') |
| 78 | valdir = os.path.join(args.data, 'val') |
| 79 | normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], |
| 80 | std=[0.229, 0.224, 0.225]) |
| 81 | transformations_train = [transforms.RandomResizedCrop(224), |
| 82 | transforms.RandomHorizontalFlip(), |
| 83 | transforms.ToTensor(), |
| 84 | normalize] |
| 85 | if args.tencrops: |
| 86 | transformations_val = [ |
| 87 | transforms.Resize(256), |
| 88 | transforms.TenCrop(224), |
| 89 | transforms.Lambda(lambda crops: torch.stack([normalize(transforms.ToTensor()(crop)) for crop in crops])), |
| 90 | ] |
| 91 | else: |
| 92 | transformations_val = [ |
| 93 | transforms.Resize(256), |
| 94 | transforms.CenterCrop(224), |
| 95 | transforms.ToTensor(), |
| 96 | normalize |
no test coverage detected