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)