def get_paras128(model_path=None):
val_trans = trans.Compose([trans.Resize(128), trans.ToTensor()])
test_data = dsets.ImageFolder(root='./dataset_pre/val/',
transform=val_trans)
test_loader = torch.utils.data.DataLoader(test_data,
batch_size=1,
shuffle=False,
num_workers=3)
model = my_model(num_classes=len(test_loader.dataset.classes),
pretrained=True)
model = torch.nn.DataParallel(model).cuda()
checkpoint = torch.load(model_path)
model.load_state_dict(checkpoint['state_dict'])
return model, test_loader
def main():
# ==================== 加载数据 ====================
global args, best_prec1
traindir = os.path.join(args.data, 'train')
valdir = os.path.join(args.data, 'val')
train_dataset = torchvision.datasets.ImageFolder(
traindir,
transforms.Compose([
transforms.Resize(128),
# transforms.RandomHorizontalFlip(),
transforms.ToTensor()
]))
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=args.batch_size, shuffle=True,
num_workers=args.workers, pin_memory=True, sampler=None)
val_loader = torch.utils.data.DataLoader(
torchvision.datasets.ImageFolder(valdir, transforms.Compose([
transforms.Resize(128),
transforms.ToTensor()
])),
batch_size=args.batch_size, shuffle=False,
num_workers=args.workers, pin_memory=True)
# ==================== 加载网络 ====================
if args.pretrained:
logger.info('==>load pretrained weights...')
model = my_model(num_classes=len(train_loader.dataset.classes), pretrained=args.pretrained)
model = torch.nn.DataParallel(model).cuda()
criterion = nn.CrossEntropyLoss(torch.cuda.FloatTensor([
1, 10.56, 4.87, 29.56, 36.45])).cuda() # 调配权重
optimizer = torch.optim.SGD(params=model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
if args.resume:
if os.path.isfile(args.resume):
logger.info("=> loading checkpoint '{0}'".format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
best_prec1 = checkpoint['best_prec1']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
logger.info("=> loaded checkpoint '{0}' (epoch {1})".format(args.resume, checkpoint['epoch']))
else:
logger.info("=> no checkpoint found at '{0}'".format(args.resume))
cudnn.benchmark = True
# ==================== 只做预测 ====================
if args.evaluate:
validate(val_loader, model, criterion)
return
# ==================== 训练网络 ====================
for epoch in range(args.start_epoch, args.epochs):
adjust_learning_rate(optimizer, epoch) # 调整学习率
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch)
# evaluate on validation set
prec1 = validate(val_loader, model, criterion)
# remember best prec@1 and save checkpoint
is_best = prec1 > best_prec1
best_prec1 = max(prec1, best_prec1)
save_checkpoint({
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
'optimizer': optimizer.state_dict(),
}, is_best)
# file_path = r'E:\Radiomics\huaxi_jiang_yinjie\cropped_images\image_roi_path_file\fixed_rec_cropped_images.xlsx'
file_path = r'E:\Radiomics\huaxi_jiang_yinjie\outcome\cross_validation\cohorts\DL_prob\DL_path\cohort5.xlsx'
cohort_num = 5
test_dataset = Yinjie_data(file_path,
transforms=valid_transforms,
train=False,
test=False)
testloader = torch.utils.data.DataLoader(test_dataset,
batch_size=8,
shuffle=False,
num_workers=0)
##加载自己更改的resnet结构,加载预训练权重
model = my_model()
pretrained_weight_path = r'E:\pycharm_project\Huaxi_Yinjie_Jiang\AP_weights\epoch_1_5.pth'
pretrained_dict = torch.load(pretrained_weight_path)
model_dict = model.state_dict()
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
if len(pretrained_dict) != 0:
print('加载成功')
# pretrained_dict = torch.load(pretrained_weight_path)
# model_dict = model.state_dict()
# pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
# model_dict.update(pretrained_dict)
# model.load_state_dict(model_dict)