def main(): parser = argparse.ArgumentParser(description='Image Classification.') parser.add_argument('--model-name', type=str, default='resnet50') parser.add_argument('--checkpoint-path',type = str, default='./checkpoints/LUNA', help= 'Path to save checkpoint, only the model with highest top1 acc will be saved,' 'And the records will also be writen in the folder') parser.add_argument('--batch-size', type = int, default=16, help='Batch size') parser.add_argument('--lr', type=float, default=1e-3, help='Initial learning rate') parser.add_argument('--epoch',type = int ,default=200, help='Maximum training epoch') parser.add_argument('--train-dir', type=str, default='./data/train', help='path to the train folder, each class has a single folder') parser.add_argument('--val-dir', type=str, default='./data/val', help='path to the validation folder, each class has a single folder' ) parser.add_argument('--test-dir', type=str, default='xxx/test', help='path to the test folder, each class has a single folder') parser.add_argument('--pretrained',type=str, default="None", help='Load which pretrained model, ' 'None : Do not load any weight, random initialize' 'Imagenet : official Imagenet pretrained model,' 'MoCo : Transfer model from Moco, path in $transfer-resume$' 'Transfer : Transfer model from Supervised pretrained, path in $transfer-resume$' 'Resume : Load checkpoint for corrupted training process, path in $resume$') parser.add_argument('--transfer-resume', type=str, default=False, help='Path to load transfering pretrained model') parser.add_argument('--resume', type = str, default='', help='Path to resume a checkpoint') parser.add_argument('--num-class', type = int, default=2, help='Number of class for the classification') parser.add_argument('--PRINT-INTERVAL', type=int, default=20, help='Number of batch to print the loss') parser.add_argument('--logpath', type=str, default='./logs/log1', help='log path') args = parser.parse_args() # tensorboard writer writer = SummaryWriter(args.logpath) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print("Device {}".format(device)) # Create checkpoint file if os.path.exists(args.checkpoint_path) == False: os.makedirs(args.checkpoint_path) normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) test_trans = transforms.Compose( [ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), normalize ] ) trainset = datasets.ImageFolder(root=args.train_dir, transform=transforms.Compose( [transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), auto.ImageNetPolicy(), transforms.ToTensor(), normalize ] )) valset = datasets.ImageFolder(root=args.val_dir, transform=test_trans ) train_loader = DataLoader(trainset, batch_size=args.batch_size, num_workers=8, sampler=ImbalancedDatasetSampler(trainset), shuffle=False) val_loader = DataLoader(valset, batch_size=args.batch_size) # Define Loss Function LOSS_FUNC = nn.CrossEntropyLoss() print(args.model_name) if args.pretrained == 'Imagenet': # ImageNet supervised pretrained model print('ImageNet supervised pretrained model') model = MODEL_DICT[args.model_name](num_classes=args.num_class, pretrained=True) elif args.pretrained == 'MoCo': # load weight from transfering model from moco print('Load weight from transfering model from moco') model = MODEL_DICT[args.model_name](num_classes=args.num_class, pretrained=False) if args.transfer_resume: if os.path.isfile(args.transfer_resume): print("=> loading checkpoint '{}'".format(args.transfer_resume)) checkpoint = torch.load(args.pretrained, map_location="cpu") # rename moco pre-trained keys state_dict = checkpoint['state_dict'] for k in list(state_dict.keys()): # retain only encoder_q up to before the embedding layer if k.startswith('module.encoder_q') and not k.startswith('module.encoder_q.fc'): # remove prefix state_dict[k[len("module.encoder_q."):]] = state_dict[k] # delete renamed or unused k del state_dict[k] msg = model.load_state_dict(state_dict, strict=False) assert set(msg.missing_keys) == {"fc.weight", "fc.bias"} print("=> loaded pre-trained model '{}'".format(args.transfer_resume)) else: print("=> no checkpoint found at '{}'".format(args.transfer_resume)) # freeze all layers but the last fc for name, param in model.named_parameters(): if name not in ['fc.weight', 'fc.bias']: param.requires_grad = False # init the fc layer model.fc.weight.data.normal_(mean=0.0, std=0.01) model.fc.bias.data.zero_() elif args.pretrained == 'Transfer': # load weight from transfering model from supervised pretraining print('Load weight from transfering model from supervised pretraining') model = MODEL_DICT[args.model_name](num_classes=args.num_class, pretrained=False) if args.transfer_resume: if os.path.isfile(args.transfer_resume): print("=> loading checkpoint '{}'".format(args.transfer_resume)) checkpoint = torch.load(args.transfer_resume) msg = model.load_state_dict(checkpoint, strict=False) assert set(msg.missing_keys) == {"fc.weight", "fc.bias"} print("=> loaded checkpoint '{}' (epoch {})" .format(args.transfer_resume, checkpoint['epoch'])) else: print("=> no checkpoint found at '{}'".format(args.transfer_resume)) for name, param in model.named_parameters(): if name not in ['fc.weight', 'fc.bias']: param.requires_grad = False # init the fc layer model.fc.weight.data.normal_(mean=0.0, std=0.01) model.fc.bias.data.zero_() else: # Random Initialize print('Random Initialize') model = MODEL_DICT[args.model_name](num_classes=args.num_class, pretrained=False) # Dataparallel for multiple GPU usage if torch.cuda.device_count() > 1: print("Let's use", torch.cuda.device_count(), "GPUs!") model = nn.DataParallel(model) model = model.to(device) # Optimizer and learning rate scheduler optimizer = torch.optim.Adam(model.parameters(),lr = args.lr) sheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10) if args.pretrained == 'Resume': # load weight from checkpoint load_resume(args, model, optimizer, args.resume) metric = [] for epoch in range(args.epoch): print("Epoch(train): ", epoch) train_loss = train(model, train_loader, optimizer, args.PRINT_INTERVAL, epoch, args, LOSS_FUNC, device) print("Epoch(val): ", epoch) # acc1, acc5, confusion_matrix, val_loss, aucs = test(model, val_loader, args.num_class, LOSS_FUNC, device) acc1, confusion_matrix, val_loss, aucs = test(model, val_loader, args.num_class, LOSS_FUNC, device) metric.append(acc1) print("Epoch Validation {}\tAcc1".format(epoch,acc1)) # Tensorboard record loss, accuracy and auc writer.add_scalars('Loss', { 'Train': train_loss, 'Val': val_loss }, epoch) writer.flush() writer.add_scalar('Accuracy', acc1, epoch) writer.flush() # Save train/val loss, acc1, acc5, confusion matrix(F1, recall, precision), AUCs record = { 'epoch': epoch + 1, 'train loss': train_loss, 'val loss': val_loss, 'acc1' : acc1, } torch.save(record, os.path.join(args.checkpoint_path, 'record epoch{}.pth.tar'.format(epoch))) # Only save the model with highest top1 acc if np.max(metric) == acc1: checkpoint = { 'epoch': epoch + 1, 'arch': args.model_name, 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict(), } torch.save(checkpoint, os.path.join(args.checkpoint_path, 'best.pth.tar')) print("Model Saved") sheduler.step(epoch) print('...........Testing..........') load_resume(args, model, optimizer, args.checkpoint_path) acc1, confusion_matrix, val_loss, aucs = test(model, val_loader, args.num_class, LOSS_FUNC, device)
def main(): parser = argparse.ArgumentParser(description='Image Classification.') parser.add_argument('--model-name', type=str, default='resnet50') parser.add_argument( '--checkpoint_path', type=str, default='/home/ubuntu/hxh/tl-ssl/tl/tl_finetune', help= 'Path to save checkpoint, only the model with highest top1 acc will be saved,' 'And the records will also be writen in the folder') parser.add_argument('--batch-size', type=int, default=128, help='Batch size') parser.add_argument('--lr', type=float, default=1e-2, help='Initial learning rate') parser.add_argument('--epoch', type=int, default=100, help='Maximum training epoch') parser.add_argument('--start-epoch', type=int, default=0, help='Start training epoch') parser.add_argument('--root-dir', type=str, default='../data/Caltech256', help='path to the image folder') parser.add_argument('--train-file', type=str, default='../dataset/caltech256_train_0.csv', help='path to the train csv file') parser.add_argument('--val-file', type=str, default='../dataset/caltech256_val_0.csv', help='path to the val csv file') parser.add_argument('--test-file', type=str, default='../dataset/caltech256_test_0.csv', help='path to the test csv file') # parser.add_argument('--test-dir', type=str, default='xxx/test', # help='path to the train folder, each class has a single folder') parser.add_argument('--cos', type=bool, default=False, help='Use cos learning rate sheduler') parser.add_argument('--schedule', default=[20, 40], nargs='*', type=int, help='learning rate schedule (when to drop lr by 10x)') parser.add_argument( '--pretrained', type=str, default='None', help='Load which pretrained model, ' 'None : Do not load any weight, random initialize' 'Imagenet : official Imagenet pretrained model,' 'MoCo : Transfer model from Moco, path in $transfer-resume$' 'Transfer : Transfer model from Supervised pretrained, path in $transfer-resume$' 'Resume : Load checkpoint for corrupted training process, path in $resume$' ) parser.add_argument( '--transfer-resume', type=str, default='/home/ubuntu/hxh/tl-ssl/tl/tl_pretrain/best.pth.tar', help='Path to load transfering pretrained model') parser.add_argument( '--resume', type=str, default='/home/ubuntu/hxh/tl-ssl/tl/tl_finetune/best.pth.tar', help='Path to resume a checkpoint') parser.add_argument('--num-class', type=int, default=257, help='Number of class for the classification') parser.add_argument('--PRINT-INTERVAL', type=int, default=20, help='Number of batch to print the loss') args = parser.parse_args() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print("Device {}".format(device)) # Create checkpoint file if os.path.exists(args.checkpoint_path) == False: os.makedirs(args.checkpoint_path) normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) test_trans = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), normalize ]) Caltech = Caltech256(root=args.root_dir, train_file=args.train_file, val_file=args.val_file, test_file=args.test_file, download=True) trainset = ImageDataset(images=Caltech.train_images, labels=Caltech.train_label, transforms=transforms.Compose([ transforms.Resize(256), transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), auto.ImageNetPolicy(), transforms.ToTensor(), normalize ])) valset = ImageDataset(images=Caltech.val_images, labels=Caltech.val_label, transforms=test_trans) testset = ImageDataset(images=Caltech.test_images, labels=Caltech.test_label, transforms=test_trans) train_loader = DataLoader(trainset, batch_size=args.batch_size, num_workers=8, shuffle=True, pin_memory=True) val_loader = DataLoader(valset, batch_size=args.batch_size, num_workers=8, pin_memory=True) test_loader = DataLoader(testset, batch_size=args.batch_size, num_workers=8, pin_memory=True) # Define Loss Function LOSS_FUNC = nn.CrossEntropyLoss().to(device) print(args.model_name) if args.pretrained == 'Imagenet': # ImageNet supervised pretrained model print('ImageNet supervised pretrained model') model = MODEL_DICT[args.model_name](num_classes=args.num_class, pretrained=True) elif args.pretrained == 'MoCo': # load weight from transfering model from moco print('Load weight from transfering model from moco') model = MODEL_DICT[args.model_name](num_classes=args.num_class, pretrained=False) if args.transfer_resume: if os.path.isfile(args.transfer_resume): print("=> loading checkpoint '{}'".format( args.transfer_resume)) checkpoint = torch.load(args.transfer_resume, map_location="cpu") # rename moco pre-trained keys state_dict = checkpoint['state_dict'] for k in list(state_dict.keys()): # retain only encoder_q up to before the embedding layer if k.startswith('module.encoder_q') and not k.startswith( 'module.encoder_q.fc'): # remove prefix state_dict[ k[len("module.encoder_q."):]] = state_dict[k] # delete renamed or unused k del state_dict[k] msg = model.load_state_dict(state_dict, strict=False) assert set(msg.missing_keys) == {"fc.weight", "fc.bias"} print("=> loaded pre-trained model '{}'".format( args.transfer_resume)) else: print("=> no checkpoint found at '{}'".format( args.transfer_resume)) # init the fc layer model.fc.weight.data.normal_(mean=0.0, std=0.01) model.fc.bias.data.zero_() elif args.pretrained == 'Transfer': # load weight from transfering model from supervised pretraining model = MODEL_DICT[args.model_name](num_classes=args.num_class, pretrained=False) print('Load weight from transfering model from supervised pretraining') if args.transfer_resume: if os.path.isfile(args.transfer_resume): print("=> loading checkpoint '{}'".format( args.transfer_resume)) checkpoint = torch.load(args.transfer_resume) state_dict = checkpoint['state_dict'] for k in list(state_dict.keys()): # retain only encoder_q up to before the embedding layer print(k) if k.startswith("module.fc.") or k.startswith("fc."): del state_dict[k] elif k.startswith('module.'): # remove prefix state_dict[k[len("module."):]] = state_dict[k] # delete renamed or unused k del state_dict[k] msg = model.load_state_dict(state_dict, strict=False) assert set(msg.missing_keys) == {"fc.weight", "fc.bias"} print("=> loaded checkpoint '{}' (epoch {})".format( args.transfer_resume, checkpoint['epoch'])) else: print("=> no checkpoint found at '{}'".format( args.transfer_resume)) # init the fc layer model.fc.weight.data.normal_(mean=0.0, std=0.01) model.fc.bias.data.zero_() else: # Random Initialize print('Random Initialize') model = MODEL_DICT[args.model_name](num_classes=args.num_class, pretrained=False) print("Let's use", torch.cuda.device_count(), "GPUs!") model = nn.DataParallel(model) model = model.to(device) # Optimizer and learning rate scheduler optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) sheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10) if args.pretrained == 'Resume': # load weight from checkpoint print('Load weight from checkpoint {}'.format(args.resume)) load_resume(args, model, optimizer, args.resume) args.start_epoch = 0 metric = [] # acc1, acc5, confusion_matrix, val_loss, aucs = test(model, test_loader, args.num_class, LOSS_FUNC, device) for epoch in range(args.start_epoch, args.epoch): # adjust_learning_rate(optimizer, epoch, args) train_loss = train(model, train_loader, optimizer, args.PRINT_INTERVAL, epoch, args, LOSS_FUNC, device) acc1, acc5, confusion_matrix, val_loss, aucs = test( model, val_loader, args.num_class, LOSS_FUNC, device) metric.append(acc1) sheduler.step() # Save train/val loss, acc1, acc5, confusion matrix(F1, recall, precision), AUCs record = { 'epoch': epoch + 1, 'train loss': train_loss, 'val loss': val_loss, 'acc1': acc1, 'acc5': acc5, 'confusion matrix': confusion_matrix, 'AUCs': aucs } torch.save( record, os.path.join(args.checkpoint_path, 'recordEpoch{}.pth.tar'.format(epoch))) # Only save the model with highest top1 acc if np.max(metric) == acc1: checkpoint = { 'epoch': epoch + 1, 'arch': args.model_name, 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict(), } torch.save(checkpoint, os.path.join(args.checkpoint_path, 'best.pth.tar')) print("Model Saved") print('...........Testing..........') load_resume(args, model, optimizer, './checkpoint/caltech_finetune/best.pth.tar') acc1, acc5, confusion_matrix, val_loss, aucs = test( model, test_loader, args.num_class, LOSS_FUNC, device)
def main(): args = parser.parse_args() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print("Device {}".format(device)) # Create checkpoint file normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) test_trans = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop((224, 224)), transforms.ToTensor(), normalize ]) trainset = DomainDataset(root_dir='path_covid/COVID-CT/Images-processed', transform=transforms.Compose([ transforms.Resize(256), transforms.RandomHorizontalFlip(), auto.ImageNetPolicy(), transforms.RandomResizedCrop(224, scale=(0.5, 1.2)), transforms.ToTensor(), normalize ])) origin_len = len(trainset) trainset, testset = torch.utils.data.random_split( trainset, [int(0.7 * len(trainset)), origin_len - int(0.7 * len(trainset))]) train_loader = DataLoader(trainset, batch_size=16, num_workers=8, shuffle=True, pin_memory=True) val_loader = DataLoader(testset, batch_size=128, num_workers=8, pin_memory=True, shuffle=True) test_loader = DataLoader(testset, batch_size=16) LOSS_FUNC = nn.CrossEntropyLoss().to(device) # Random Initialize print('Random Initialize') model = simple_classfifer() # model = torchvision.models.resnet50(num_classes=2, pretrained=False) # Dataparallel for multiple GPU usage if torch.cuda.device_count() > 1: print("Let's use", torch.cuda.device_count(), "GPUs!") model = nn.DataParallel(model) model = model.to(device) # Optimizer and learning rate scheduler optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9) # optimizer = torch.optim.RMSprop(model.parameters(), lr=args.lr) # print('start test') # for index, (images, target) in enumerate(tqdm(val_loader)): # print(target) total_epoch = 10 metric = [] acc1, acc5, confusion_matrix, val_loss, aucs = test( model, val_loader, 2, LOSS_FUNC, device) for epoch in range(1, total_epoch): train_loss = train(model, train_loader, optimizer, 3, epoch, LOSS_FUNC, device) acc1, acc5, confusion_matrix, val_loss, aucs = test( model, val_loader, 2, LOSS_FUNC, device) metric.append(acc1) # Save train/val loss, acc1, acc5, confusion matrix(F1, recall, precision), AUCs record = { 'epoch': epoch + 1, 'train loss': train_loss, 'val loss': val_loss, 'acc1': acc1, 'acc5': acc5, 'confusion matrix': confusion_matrix, 'AUCs': aucs } # torch.save(record, os.path.join('./','recordEpoch{}.pth.tar'.format(epoch))) # Only save the model with highest top1 acc if np.max(metric) == acc1: checkpoint = { 'epoch': epoch + 1, 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict(), } torch.save(checkpoint, os.path.join('./', 'nih_flo_best.pth.tar')) print("Model Saved") print('...........Testing..........') load_resume(args, model, optimizer, os.path.join('./', 'nih_flo_best.pth.tar')) acc1, acc5, confusion_matrix, val_loss, aucs = test( model, test_loader, 2, LOSS_FUNC, device)