def run(criterion, net_name, use_dataset_index, use_mask, save_img_pred: bool):
path_data = config.path_data[use_dataset_index]
# 文件路径
path_data_root = path_data["dataset"]
path_checkpoints = path_data["checkpoints"]
path_result = os.path.join(path_data["result"], net_name)
os.makedirs(path_result, exist_ok=True)
net = UNet().to(config.device)
for param in net.parameters():
param.requires_grad = False
path_pert = "/home/pengzx/deepLearning/result/Glaucoma/UNet/pert/"
iou_total = []
for index in config.checkpoints_indexes:
path_pert_save = os.path.join(path_pert, str(index))
attacked_dataset = MyAttackedDataset(path_data_root=path_data_root,
phase="train",
path_pert=path_pert_save,
transform_list=transform_compose)
attacked_data_loader = DataLoader(attacked_dataset,
batch_size=1,
shuffle=False)
for index_2 in config.checkpoints_indexes:
net.load_state_dict(
torch.load(os.path.join(path_checkpoints,
"{}_{}.pth").format(net_name, index_2),
map_location=config.device))
iou_list = []
for i, (img, label, mask, pert,
name) in enumerate(attacked_data_loader):
img, pert = img.to(config.device), pert.to(config.device)
img_pert = img.clone() + pert
pred = net(img_pert)
pred[pred > 0] = 1.
pred[pred < 1.] = 0.
label[label > 0.5] = 1.
label[label < 1.] = 0.
iou = get_iou(pred[0].data.cpu().numpy(),
label.data.cpu().numpy())
iou_list.append(iou)
iou_list = np.array(iou_list)
print("模型[{}]产生的扰动,模型[{}]的预测结果iou={}", index, index_2,
iou_list.mean())
iou_total.append(iou_list)
iou_total = np.array(iou_total)
np.save("./iou_total.npy", iou_total)
def build_model(model_name, num_classes):
if model_name == 'SQNet':
return SQNet(classes=num_classes)
elif model_name == 'LinkNet':
return LinkNet(classes=num_classes)
elif model_name == 'SegNet':
return SegNet(classes=num_classes)
elif model_name == 'UNet':
return UNet(classes=num_classes)
elif model_name == 'ENet':
return ENet(classes=num_classes)
elif model_name == 'ERFNet':
return ERFNet(classes=num_classes)
elif model_name == 'CGNet':
return CGNet(classes=num_classes)
elif model_name == 'EDANet':
return EDANet(classes=num_classes)
elif model_name == 'ESNet':
return ESNet(classes=num_classes)
elif model_name == 'ESPNet':
return ESPNet(classes=num_classes)
elif model_name == 'LEDNet':
return LEDNet(classes=num_classes)
elif model_name == 'ESPNet_v2':
return EESPNet_Seg(classes=num_classes)
elif model_name == 'ContextNet':
return ContextNet(classes=num_classes)
elif model_name == 'FastSCNN':
return FastSCNN(classes=num_classes)
elif model_name == 'DABNet':
return DABNet(classes=num_classes)
elif model_name == 'FSSNet':
return FSSNet(classes=num_classes)
elif model_name == 'FPENet':
return FPENet(classes=num_classes)
def create_refined_unet_v2(input_channels, num_classes, radius=20, eps=1e-4, theta_gamma=3.0, spatial_compat=3.0, bilateral_compat=10.0, num_iterations=10, gt_prob=0.7, unet_pretrained=None):
""" Create Refined UNet v2 """
# Input
inputs = tf.keras.Input(
shape=[None, None, input_channels], name='inputs')
# Create UNet
unet = UNet()
# Restore pretrained UNet
if unet_pretrained:
checkpoint = tf.train.Checkpoint(model=unet)
checkpoint.restore(tf.train.latest_checkpoint(unet_pretrained))
print('Checkpoint restored, at {}'.format(
tf.train.latest_checkpoint(unet_pretrained)))
# Create CRF layer
crf = CRFLayer(num_classes, radius, eps, theta_gamma,
spatial_compat, bilateral_compat, num_iterations)
# RGB channels
image = inputs[..., 4:1:-1] * 255
# Only forward
logits = unet(inputs)
probs = tf.nn.softmax(logits, name='logits2probs')
unary = -tf.math.log(probs * gt_prob, name='probs2unary')
refined_logits = crf(unary=unary, image=image)
return tf.keras.Model(inputs=inputs, outputs=[logits, refined_logits])
def build_model(model_name, num_classes):
# small model
if model_name == 'ENet':
return ENet(classes=num_classes)
elif model_name == 'ERFNet':
return ERFNet(classes=num_classes)
elif model_name == 'ESPNet':
return ESPNet(classes=num_classes)
elif model_name == 'ESPNet_v2':
return EESPNet_Seg(classes=num_classes)
elif model_name == 'DABNet':
return DABNet(classes=num_classes)
elif model_name == 'BiSeNetV2':
return BiSeNetV2(n_classes=num_classes)
# large model
elif model_name == 'UNet':
return UNet(classes=num_classes)
elif model_name == 'PSPNet50':
return PSPNet(layers=50,
bins=(1, 2, 3, 6),
dropout=0.1,
classes=num_classes,
zoom_factor=8,
use_ppm=True,
pretrained=True)
# elif model_name == 'PSANet50':
# return PSANet(layers=50, dropout=0.1, classes=num_classes, zoom_factor=8, use_psa=True, psa_type=2, compact=compact,
# shrink_factor=shrink_factor, mask_h=mask_h, mask_w=mask_w, psa_softmax=True, pretrained=True)
elif model_name == 'Deeplabv3plus':
return Deeplabv3plus(cfg, num_classes=num_classes)
def get_net(net_name):
if net_name == "UNet":
return UNet(3, 1)
if net_name == "CENet":
return CENet()
if net_name == "FCN":
vgg_model = VGGNet(requires_grad=True, show_params=False)
return FCNs(pretrained_net=vgg_model, n_class=1)
return None
def train():
ex = wandb.init(project="PQRST-segmentation")
ex.config.setdefaults(wandb_config)
logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
logging.info(f'Using device {device}')
net = UNet(in_ch=1, out_ch=4)
net.to(device)
try:
train_model(net=net, device=device, batch_size=wandb.config.batch_size, lr=wandb.config.lr, epochs=wandb.config.epochs)
except KeyboardInterrupt:
try:
save = input("save?(y/n)")
if save == "y":
torch.save(net.state_dict(), 'net_params.pkl')
sys.exit(0)
except SystemExit:
os._exit(0)
def main(args):
if not os.path.exists(args.save_path):
os.mkdir(args.save_path)
net = UNet(n_channels=3, n_classes=1)
checkpoint = flow.load(args.pretrained_path)
net.load_state_dict(checkpoint)
net.to("cuda")
x_test_dir, y_test_dir = get_datadir_path(args, split="test")
test_dataset = Dataset(
x_test_dir, y_test_dir, augmentation=get_test_augmentation(),
)
print("Begin Testing...")
for i, (image, mask) in enumerate(tqdm(test_dataset)):
show_image = image
with flow.no_grad():
image = image / 255.0
image = image.astype(np.float32)
image = flow.tensor(image, dtype=flow.float32)
image = image.permute(2, 0, 1)
image = image.to("cuda")
pred = net(image.unsqueeze(0).to("cuda"))
pred = pred.numpy()
pred = pred > 0.5
save_picture_name = os.path.join(args.save_path, "test_image_" + str(i))
visualize(
save_picture_name, image=show_image, GT=mask[0, :, :], Pred=pred[0, 0, :, :]
)
def main(config):
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
train_dataset, valid_dataset = generate_datasets(
config['data_dir'], valid_ids=config['val_ids'])
# TODO: define and add data augmentation + image normalization
# train_dataset.transform = train_transform
# valid_dataset.transform = valid_transform
transforms = A.Compose([
A.Normalize(), # TODO: change values
ToTensorV2()
])
train_dataset.transform = transforms
valid_dataset.transform = transforms
train_loader = DataLoader(train_dataset,
batch_size=config['batch_size'],
shuffle=True,
num_workers=config['num_workers'])
valid_loader = DataLoader(valid_dataset,
config['batch_size'],
shuffle=False,
num_workers=config['num_workers'])
model = UNet()
model = model.to(device)
criterion = config['criterion']
optimizer = torch.optim.Adam(params=model.parameters(), lr=3e-4)
trainer = Trainer(model=model,
criterion=criterion,
optimizer=optimizer,
config=config,
train_loader=train_loader,
val_loader=valid_loader,
device=device)
trainer.train()
return model
def build_model(model_name, num_classes):
# for deeplabv3
model_map = {
'deeplabv3_resnet50': network.deeplabv3_resnet50,
'deeplabv3plus_resnet50': network.deeplabv3plus_resnet50,
'deeplabv3_resnet101': network.deeplabv3_resnet101,
'deeplabv3plus_resnet101': network.deeplabv3plus_resnet101,
'deeplabv3_mobilenet': network.deeplabv3_mobilenet,
'deeplabv3plus_mobilenet': network.deeplabv3plus_mobilenet
}
if model_name == 'SQNet':
return SQNet(classes=num_classes)
elif model_name == 'LinkNet':
return LinkNet(classes=num_classes)
elif model_name == 'SegNet':
return SegNet(classes=num_classes)
elif model_name == 'UNet':
return UNet(classes=num_classes)
elif model_name == 'ENet':
return ENet(classes=num_classes)
elif model_name == 'ERFNet':
return ERFNet(classes=num_classes)
elif model_name == 'CGNet':
return CGNet(classes=num_classes)
elif model_name == 'EDANet':
return EDANet(classes=num_classes)
elif model_name == 'ESNet':
return ESNet(classes=num_classes)
elif model_name == 'ESPNet':
return ESPNet(classes=num_classes)
elif model_name == 'LEDNet':
return LEDNet(classes=num_classes)
elif model_name == 'ESPNet_v2':
return EESPNet_Seg(classes=num_classes)
elif model_name == 'ContextNet':
return ContextNet(classes=num_classes)
elif model_name == 'FastSCNN':
return FastSCNN(classes=num_classes)
elif model_name == 'DABNet':
return DABNet(classes=num_classes)
elif model_name == 'FSSNet':
return FSSNet(classes=num_classes)
elif model_name == 'FPENet':
return FPENet(classes=num_classes)
elif model_name == 'FCN':
return FCN32VGG(classes=num_classes)
elif model_name in model_map.keys():
return model_map[model_name](num_classes, output_stride=8)
def __model(self, tf_mix):
# define model flow
# stft
stft_module = STFT_Module(
frame_length=self.stft_params["frame_length"],
frame_step=self.stft_params["frame_step"],
fft_length=self.stft_params["fft_length"],
pad_end=self.stft_params["pad_end"],
epsilon=self.epsilon)
# mix data transform
tf_spec_mix = stft_module.STFT(tf_mix)
print("spec mix", tf_spec_mix.dtype)
tf_spec_mix = stft_module.to_T_256(
tf_spec_mix) # cut time dimension to 256 for u-net architecture
tf_phase_mix = tf.sign(tf_spec_mix)
tf_phase_mix = self.expand_channel(tf_phase_mix)
# tf_mag_spec_mix = stft_module.to_magnitude_spec(tf_spec_mix, normalize=False)
tf_amp_spec_mix = stft_module.to_amp_spec(tf_spec_mix, normalize=False)
tf_mag_spec_mix = tf.log(tf_amp_spec_mix + self.epsilon)
tf_mag_spec_mix = tf.expand_dims(tf_mag_spec_mix,
-1) # (Batch, Time, Freq, Channel))
tf_amp_spec_mix = tf.expand_dims(tf_amp_spec_mix, -1)
tf_f_512_mag_spec_mix = stft_module.to_F_512(tf_mag_spec_mix)
# target data transform
# tf_spec_target = stft_module.STFT(tf_target)
# tf_spec_target = stft_module.to_T_256(tf_spec_target) # cut time dimensiton to 256 for u-net architecture
# tf_amp_spec_target = stft_module.to_amp_spec(tf_spec_target, normalize=False)
# tf_amp_spec_target = tf.expand_dims(tf_amp_spec_target, -1)
u_net = UNet(input_shape=(tf_f_512_mag_spec_mix.shape[1:]))
tf_est_masks = u_net(tf_f_512_mag_spec_mix)
#F: 512 → 513
zero_pad = tf.zeros_like(tf_mag_spec_mix)
zero_pad = tf.expand_dims(zero_pad[:, :, 1, :], -1)
tf_est_masks = tf.concat([tf_est_masks, zero_pad], 2)
tf_est_spec = tf.math.multiply(tf_est_masks, tf_amp_spec_mix)
tf_est_source_spec = tf.math.multiply(tf.complex(tf_est_spec, 0.),
tf_phase_mix)
tf_est_source_spec = tf.squeeze(tf_est_source_spec, axis=-1)
est_source = stft_module.ISTFT(tf_est_source_spec)
return est_source
def __model(self, tf_mix, tf_target, tf_lr):
# define model flow
# stft
stft_module = STFT_Module(
frame_length = self.stft_params["frame_length"],
frame_step= self.stft_params["frame_step"],
fft_length = self.stft_params["fft_length"],
epsilon = self.epsilon,
pad_end = self.stft_params["pad_end"]
)
# mix data transform
tf_spec_mix = stft_module.STFT(tf_mix)
# tf_mag_spec_mix = stft_module.to_magnitude_spec(tf_spec_mix, normalize=False)
tf_amp_spec_mix = stft_module.to_amp_spec(tf_spec_mix, normalize =False)
tf_mag_spec_mix = tf.log(tf_amp_spec_mix + self.epsilon)
tf_mag_spec_mix = tf.expand_dims(tf_mag_spec_mix, -1)# (Batch, Time, Freq, Channel))
tf_amp_spec_mix = tf.expand_dims(tf_amp_spec_mix, -1)
tf_f_512_mag_spec_mix = stft_module.to_F_512(tf_mag_spec_mix)
# target data transform
tf_spec_target = stft_module.STFT(tf_target)
tf_amp_spec_target = stft_module.to_amp_spec(tf_spec_target, normalize=False)
tf_amp_spec_target = tf.expand_dims(tf_amp_spec_target, -1)
u_net = UNet(
input_shape =(
tf_f_512_mag_spec_mix.shape[1:]
)
)
tf_est_masks = u_net(tf_f_512_mag_spec_mix)
#F: 512 → 513
zero_pad = tf.zeros_like(tf_mag_spec_mix)
zero_pad = tf.expand_dims(zero_pad[:,:,1,:], -1)
tf_est_masks = tf.concat( [tf_est_masks, zero_pad], 2)
# tf_est_spec = tf.math.multiply(tf_est_masks, tf_amp_spec_mix)
tf_ora_masks = Masks.iaf(tf_amp_spec_mix, tf_amp_spec_target,self.epsilon)
tf_loss = 10 * Loss.mean_square_error(tf_est_masks, tf_ora_masks)
tf_train_step = Trainer.Adam(tf_loss, tf_lr)
return tf_train_step, tf_loss, tf_amp_spec_target, tf_mag_spec_mix, tf_spec_mix, tf_est_masks, tf_ora_masks
def __model(self, tf_mix):
# define model flow
# stft
stft_module = STFT_Module(
frame_length = self.stft_params["frame_length"],
frame_step= self.stft_params["frame_step"],
fft_length = self.stft_params["fft_length"],
epsilon = self.epsilon,
pad_end = self.stft_params["pad_end"]
)
# mix data transform
tf_spec_mix = stft_module.STFT(tf_mix)
tf_phase_mix = tf.sign(tf_spec_mix)
tf_phase_mix = self.expand_channel(tf_phase_mix)
tf_amp_spec_mix = stft_module.to_amp_spec(tf_spec_mix, normalize =False)
tf_mag_spec_mix = tf.log(tf_amp_spec_mix + self.epsilon)
tf_mag_spec_mix = tf.expand_dims(tf_mag_spec_mix, -1)# (Batch, Time, Freq, Channel))
tf_amp_spec_mix = tf.expand_dims(tf_amp_spec_mix, -1)
tf_mag_spec_mix = tf_mag_spec_mix[:,:,:1024,:]
u_net = UNet(
input_shape =(
tf_mag_spec_mix.shape[1:]
)
)
tf_est_masks = u_net(tf_mag_spec_mix)
zero_pad = tf.zeros_like(tf_mag_spec_mix)
zero_pad = tf.expand_dims(zero_pad[:,:,1,:], -1)
tf_est_masks = tf.concat([zero_pad, tf_est_masks], 2)
tf_est_spec = tf.math.multiply(tf_est_masks, tf_amp_spec_mix)
tf_est_source_spec = tf.math.multiply(tf.complex(tf_est_spec, 0.), tf_phase_mix)
tf_est_source_spec = tf.squeeze(tf_est_source_spec, axis=-1)
est_source = stft_module.ISTFT(tf_est_source_spec)
return est_source
def __model(self, tf_mix, tf_target, tf_lr):
# define model flow
# stft
stft_module = STFT_Module(
frame_length=self.stft_params["frame_length"],
frame_step=self.stft_params["frame_step"],
fft_length=self.stft_params["fft_length"],
epsilon=self.epsilon,
pad_end=self.stft_params["pad_end"])
mr1_stft_module = STFT_Module(
frame_length=self.mr1_stft_params["frame_length"],
frame_step=self.mr1_stft_params["frame_step"],
fft_length=self.mr1_stft_params["fft_length"],
epsilon=self.epsilon,
pad_end=self.mr1_stft_params["pad_end"])
mr2_stft_module = STFT_Module(
frame_length=self.mr2_stft_params["frame_length"],
frame_step=self.mr2_stft_params["frame_step"],
fft_length=self.mr2_stft_params["fft_length"],
epsilon=self.epsilon,
pad_end=self.mr2_stft_params["pad_end"])
# print(tf_mix.shape)
# tf_mix = stft_module.zero_padding(tf_mix, self.sample_len, self.train_data_num)
# print(tf_mix.shape)
# mix data transform
tf_spec_mix = stft_module.STFT(tf_mix)
tf_amp_spec_mix = stft_module.to_amp_spec(tf_spec_mix, normalize=False)
tf_mag_spec_mix = tf.log(tf_amp_spec_mix + self.epsilon)
tf_mag_spec_mix = tf.expand_dims(tf_mag_spec_mix,
-1) # (Batch, Time, Freq, Channel))
tf_amp_spec_mix = tf.expand_dims(tf_amp_spec_mix, -1)
#mr1 mix data transform
tf_mr1_spec_mix = mr1_stft_module.STFT(tf_mix)
tf_mr1_amp_spec_mix = stft_module.to_amp_spec(tf_mr1_spec_mix,
normalize=False)
tf_mr1_mag_spec_mix = tf.log(tf_mr1_amp_spec_mix + self.epsilon)
tf_mr1_mag_spec_mix = tf.expand_dims(
tf_mr1_mag_spec_mix, -1) # (Batch, Time, Freq, Channel))
tf_mr2_spec_mix = mr2_stft_module.STFT(tf_mix)
tf_mr2_amp_spec_mix = stft_module.to_amp_spec(tf_mr2_spec_mix,
normalize=False)
tf_mr2_mag_spec_mix = tf.log(tf_mr2_amp_spec_mix + self.epsilon)
tf_mr2_mag_spec_mix = tf.expand_dims(tf_mr2_mag_spec_mix, -1)
# # target data transform
tf_spec_target = stft_module.STFT(tf_target)
tf_amp_spec_target = stft_module.to_amp_spec(tf_spec_target,
normalize=False)
tf_amp_spec_target = tf.expand_dims(tf_amp_spec_target, -1)
tf_input_spec = tf.concat(
[tf_mag_spec_mix, tf_mr1_mag_spec_mix, tf_mr2_mag_spec_mix], 3)
print(tf_input_spec.shape)
tf_input_spec = tf_input_spec[:, :, :1024, :]
print(tf_input_spec.shape)
u_net = UNet(input_shape=tf_input_spec.shape[1:])
tf_est_masks = u_net(tf_input_spec)
zero_pad = tf.zeros_like(tf_mag_spec_mix)
zero_pad = tf.expand_dims(zero_pad[:, :, 1, :], -1)
tf_est_masks = tf.concat([zero_pad, tf_est_masks], 2)
tf_est_spec = tf.math.multiply(tf_est_masks, tf_amp_spec_mix)
tf_loss = 10 * Loss.mean_square_error(tf_est_spec, tf_amp_spec_target)
tf_train_step = Trainer.Adam(tf_loss, tf_lr)
return tf_train_step, tf_loss, tf_amp_spec_target, tf_mag_spec_mix, tf_spec_mix, tf_est_masks, tf_est_spec
import torch
from model.UNet import UNet
torch.manual_seed(0)
if __name__ == "__main__":
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Device being used: {}".format(DEVICE))
# UNet
u_net = UNet(in_channels=1, out_channels=2, hidden_channels=64).to(DEVICE)
print(u_net)
random_test = torch.randn((1, 1, 572, 572)).to(DEVICE)
out = u_net(random_test)
print("Output is of size: {}".format(out.size()))
del u_net
print("Program has Ended")
def build_model(model_name,
num_classes,
backbone='resnet18',
pretrained=False,
out_stride=32,
mult_grid=False):
if model_name == 'FCN':
model = FCN(num_classes=num_classes)
elif model_name == 'FCN_ResNet':
model = FCN_ResNet(num_classes=num_classes,
backbone=backbone,
out_stride=out_stride,
mult_grid=mult_grid)
elif model_name == 'SegNet':
model = SegNet(classes=num_classes)
elif model_name == 'UNet':
model = UNet(num_classes=num_classes)
elif model_name == 'BiSeNet':
model = BiSeNet(num_classes=num_classes, backbone=backbone)
elif model_name == 'BiSeNetV2':
model = BiSeNetV2(num_classes=num_classes)
elif model_name == 'HRNet':
model = HighResolutionNet(num_classes=num_classes)
elif model_name == 'Deeplabv3plus_res101':
model = DeepLabv3_plus(nInputChannels=3,
n_classes=num_classes,
os=out_stride,
pretrained=True)
elif model_name == "DDRNet":
model = DDRNet(pretrained=True, num_classes=num_classes)
elif model_name == 'PSPNet_res50':
model = PSPNet(layers=50,
bins=(1, 2, 3, 6),
dropout=0.1,
num_classes=num_classes,
zoom_factor=8,
use_ppm=True,
pretrained=True)
elif model_name == 'PSPNet_res101':
model = PSPNet(layers=101,
bins=(1, 2, 3, 6),
dropout=0.1,
num_classes=num_classes,
zoom_factor=8,
use_ppm=True,
pretrained=True)
# elif model_name == 'PSANet50':
# return PSANet(layers=50, dropout=0.1, classes=num_classes, zoom_factor=8, use_psa=True, psa_type=2, compact=compact,
# shrink_factor=shrink_factor, mask_h=mask_h, mask_w=mask_w, psa_softmax=True, pretrained=True)
if pretrained:
checkpoint = model_zoo.load_url(model_urls[backbone])
model_dict = model.state_dict()
# print(model_dict)
# 筛除不加载的层结构
pretrained_dict = {
'backbone.' + k: v
for k, v in checkpoint.items() if 'backbone.' + k in model_dict
}
# 更新当前网络的结构字典
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
return model
def main():
# 参数
args = get_args()
if not osp.exists(args.result_dir):
os.makedirs(args.result_dir)
print("Evaluating configuration:")
for arg in vars(args):
print("{}:\t{}".format(arg, getattr(args, arg)))
with open('eval-config.json', 'w') as f:
json.dump(args.__dict__, f, indent=4)
# 数据
if args.test:
dataset = SpineDataset(root=args.root,
split='test',
transform=test_transform)
else:
dataset = SpineDataset(root=args.root,
split='val',
transform=val_transform)
dataloader = DataLoader(dataset,
batch_size=1,
shuffle=False,
num_workers=4,
pin_memory=True)
# 模型
os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda
if torch.cuda.is_available():
device = torch.device('cuda')
else:
device = torch.device('cpu')
if args.network == 'DeepLab':
model = gcv.models.DeepLabV3(nclass=args.num_classes,
backbone=args.backbone)
elif args.network == 'FCN':
model = gcv.models.FCN(nclass=args.num_classes, backbone=args.backbone)
elif args.network == 'PSPNet':
model = gcv.models.PSP(nclass=args.num_classes, backbone=args.backbone)
elif args.network == 'UNet':
model = UNet(n_class=args.num_classes, backbone=args.backbone)
print('load model from {} ...'.format(args.model))
model.load_state_dict(
torch.load(args.model, map_location='cpu')['state_dict'])
model = model.to(device)
print('Done!')
# 测试
def eval():
with torch.no_grad():
model.eval()
result = []
tq = tqdm.tqdm(total=len(dataloader))
if args.test:
tq.set_description('test')
for i, (data, img_file) in enumerate(dataloader):
tq.update(1)
data = data.to(device)
predict = np.zeros(
(data.size()[1], data.size()[3], data.size()[4]),
dtype=np.uint16)
for idx in range(data.size()[1]):
if args.network in ['DeepLab', 'FCN', 'PSPNet']:
final_out = model(data[:, idx])[0]
elif args.network == 'UNet':
final_out = model(data[:, idx])
predict[idx] = final_out.argmax(
dim=1).cpu().squeeze().numpy().astype(np.uint16)
pred_img = sitk.GetImageFromArray(predict)
test_img = sitk.ReadImage(
osp.join(args.root, 'test', 'image', img_file[0]))
pred_img.CopyInformation(test_img)
result_file = 'mask_' + img_file[0].lower()
sitk.WriteImage(pred_img,
osp.join(args.result_dir, result_file))
else:
tq.set_description('val')
for i, (data, mask, mask_file) in enumerate(dataloader):
tq.update(1)
gt_img = sitk.ReadImage(
osp.join(args.root, 'val', 'groundtruth',
mask_file[0]))
data = data.to(device)
predict = np.zeros(
(data.size()[1], data.size()[3], data.size()[4]),
dtype=np.uint16)
for idx in range(data.size()[1]):
if args.network in ['DeepLab', 'FCN', 'PSPNet']:
final_out = model(data[:, idx])[0]
elif args.network == 'UNet':
final_out = model(data[:, idx])
predict[idx] = final_out.argmax(
dim=1).cpu().squeeze().numpy().astype(np.uint16)
pred_img = sitk.GetImageFromArray(predict)
pred_img.CopyInformation(gt_img)
sitk.WriteImage(pred_img,
osp.join(args.result_dir, mask_file[0]))
ppv, sensitivity, dice, _ = metrics.precision_recall_fscore_support(
mask.numpy().flatten(),
predict.flatten(),
average='binary')
result.append([dice, ppv, sensitivity])
result = np.array(result)
result_mean = result.mean(axis=0)
result_std = result.std(axis=0)
print(result_mean, result_std)
np.savetxt(osp.join(args.result_dir, 'result.txt'),
result_mean,
fmt='%.3f',
header='Dice, Sensitivity, PPV')
tq.close()
eval()
def build_model(model_name, num_classes):
if model_name == 'DABNet':
return DABNet(classes=num_classes)
elif model_name == 'FCN_8S_res18':
return FCN_res(backbone='resnet18', classes=num_classes, pretrained=True, scale=8)
elif model_name == 'FCN_8S_res34':
return FCN_res(backbone='resnet34', classes=num_classes, pretrained=True, scale=8)
elif model_name == 'FCN_8S_res50':
return FCN_res(backbone='resnet50', classes=num_classes, pretrained=True, scale=8)
elif model_name == 'FCN_8S_res101':
return FCN_res(backbone='resnet101', classes=num_classes, pretrained=True, scale=8)
elif model_name == 'FCN_32S_res18':
return FCN_res(backbone='resnet18', classes=num_classes, pretrained=True, scale=32)
elif model_name == 'FCN_32S_res50':
return FCN_res(backbone='resnet50', classes=num_classes, pretrained=True, scale=32)
elif model_name == 'FCN_32S_res101':
return FCN_res(backbone='resnet101', classes=num_classes, pretrained=True, scale=32)
elif model_name == 'UNet_res18':
return UNet_res(backbone='resnet18', pretrained=True, classes=num_classes)
elif model_name == 'UNet_res34':
return UNet_res(backbone='resnet34', pretrained=True, classes=num_classes)
elif model_name == 'UNet_res50':
return UNet_res(backbone='resnet50', pretrained=True, classes=num_classes)
elif model_name == 'UNet_res101':
return UNet_res(backbone='resnet101', pretrained=True, classes=num_classes)
elif model_name == 'UNet_res18_ori':
return UNet_res_ori(backbone='resnet18', pretrained=True, classes=num_classes)
elif model_name == 'UNet_res34_ori':
return UNet_res_ori(backbone='resnet34', pretrained=True, classes=num_classes)
elif model_name == 'UNet_res50_ori':
return UNet_res_ori(backbone='resnet50', pretrained=True, classes=num_classes)
elif model_name == 'UNet_res101_ori':
return UNet_res_ori(backbone='resnet101', pretrained=True, classes=num_classes)
elif model_name == 'PSPNet_res18':
return PSPNet(layers=18, bins=(1, 2, 3, 6), dropout=0.1, classes=num_classes, zoom_factor=8, use_ppm=True,
pretrained=True)
elif model_name == 'PSPNet_res34':
return PSPNet(layers=34, bins=(1, 2, 3, 6), dropout=0.1, classes=num_classes, zoom_factor=8, use_ppm=True,
pretrained=True)
elif model_name == 'PSPNet_res50':
return PSPNet(layers=50, bins=(1, 2, 3, 6), dropout=0.1, classes=num_classes, zoom_factor=8, use_ppm=True,
pretrained=True)
elif model_name == 'PSPNet_res101':
return PSPNet(layers=101, bins=(1, 2, 3, 6), dropout=0.1, classes=num_classes, zoom_factor=8, use_ppm=True,
pretrained=True)
## backbone == vgg
elif model_name == 'UNet':
return UNet(classes=num_classes)
elif model_name == 'UNet_overlap':
return UNet_overlap(classes=num_classes)
elif model_name == 'BiSeNet_res18':
return BiSeNet(backbone='resnet18', n_classes=num_classes, pretrained=False)
elif model_name == 'BiSeNet_res101':
return BiSeNet(backbone='resnet101', n_classes=num_classes, pretrained=False)
elif model_name == 'lightSeg':
return lightSeg(backbone='resnet101', n_classes=num_classes, pretrained=False)
elif model_name == 'ENet':
return ENet(classes=num_classes)
## GALDNet
elif model_name == 'GALD_res50':
return GALD_res50(num_classes=num_classes)
def main():
""" DataLoader """
train_data = PartAffordanceDataset(
'train.csv',
transform=transforms.Compose([CenterCrop(),
ToTensor(),
Normalize()]))
test_data = PartAffordanceDataset(
'test.csv',
transform=transforms.Compose([CenterCrop(),
ToTensor(),
Normalize()]))
train_loader = DataLoader(train_data,
batch_size=args.batch_size,
shuffle=True,
num_workers=4)
test_loader = DataLoader(test_data,
batch_size=args.batch_size,
shuffle=False,
num_workers=4)
if args.model == 'FCN8s':
model = FCN8s(args.in_channel, args.n_classes)
elif args.model == 'SegNetBasic':
model = SegNetBasic(args.in_channel, args.n_classes)
elif args.model == 'UNet':
model = UNet(args.in_channel, args.n_classes)
else:
print('This model doesn\'t exist in the model directory')
sys.exit(1)
model.apply(init_weight)
""" training """
if args.writer:
writer = SummaryWriter(args.result_path)
if args.class_weight:
criterion = nn.CrossEntropyLoss(weight=class_weight.to(args.device))
else:
criterion = nn.CrossEntropyLoss()
model.to(args.device)
optimizer = optim.Adam(model.parameters(), lr=args.learning_rate)
train_losses = []
val_iou = []
mean_iou = []
best_mean_iou = 0.0
for epoch in range(args.max_epoch):
model.train()
running_loss = 0.0
for i, sample in tqdm.tqdm(enumerate(train_loader),
total=len(train_loader)):
optimizer.zero_grad()
x, y = sample['image'], sample['class']
x = x.to(args.device)
y = y.to(args.device)
h = model(x)
loss = criterion(h, y)
loss.backward()
optimizer.step()
running_loss += loss.item()
train_losses.append(running_loss / i)
val_iou.append(
eval_model(model, test_loader, args.device).to('cpu').float())
mean_iou.append(val_iou[-1].mean().item())
if best_mean_iou < mean_iou[-1]:
best_mean_iou = mean_iou[-1]
torch.save(model.state_dict(),
args.result_path + '/best_mean_iou_model.prm')
if writer is not None:
writer.add_scalar("train_loss", train_losses[-1], epoch)
writer.add_scalar("mean_IoU", mean_iou[-1], epoch)
writer.add_scalars(
"class_IoU", {
'iou of class 0': val_iou[-1][0],
'iou of class 1': val_iou[-1][1],
'iou of class 2': val_iou[-1][2],
'iou of class 3': val_iou[-1][3],
'iou of class 4': val_iou[-1][4],
'iou of class 5': val_iou[-1][5],
'iou of class 6': val_iou[-1][6],
'iou of class 7': val_iou[-1][7]
}, epoch)
print('epoch: {}\tloss: {:.5f}\tmean IOU: {:.3f}'.format(
epoch, train_losses[-1], mean_iou[-1]))
torch.save(model.state_dict(), args.result_path + "/final_model.prm")
# must be a multiple of 32
class CenterCrop(object):
def __call__(self, sample):
image, cls = sample['image'], sample['class']
image = Image.fromarray(np.uint8(image))
image = crop_center_pil_image(image, 320, 256)
cls = crop_center_numpy(cls, 256, 320)
image = np.asarray(image)
return {'image': image, 'class': cls}
elif args.model == 'SegNetBasic':
model = SegNetBasic(args.in_channel, args.n_classes)
elif args.model == 'UNet':
model = UNet(args.in_channel, args.n_classes)
else:
print('This model doesn\'t exist in the model directory')
sys.exit(1)
if args.params_path is not None:
model.load_state_dict(
torch.load(args.params_path,
map_location=lambda storage, loc: storage))
""" define DataLoader """
data = PartAffordanceDataset('test.csv',
transform=transforms.Compose(
[CenterCrop(),
ToTensor(),
Normalize()]))
def train(fold_idx=1):
# 1. Load dataset
dataset_train = ICH_CT_32(
ROOT=config['dataset_root'],
transform=T.Compose([T.ToTensor(),
T.Normalize([
0.5,
], [
0.5,
])]),
is_train=True,
fold_idx=fold_idx)
dataloader_train = DataLoader(dataset_train,
batch_size=config['batch_size'],
shuffle=True,
num_workers=1)
dataset_eval = ICH_CT_32(ROOT=config['dataset_root'],
transform=T.Compose(
[T.ToTensor(),
T.Normalize([
0.5,
], [
0.5,
])]),
is_train=False,
fold_idx=fold_idx)
dataloader_eval = DataLoader(dataset_eval,
batch_size=config['batch_size'],
shuffle=False,
num_workers=1)
# 2. Build model
net = UNet()
# net.finetune_from('pretrained_weights/vgg16-397923af.pth')
net = nn.DataParallel(net, device_ids=[0])
print(net)
# 3. Criterion
criterion = nn.CrossEntropyLoss(weight=torch.tensor([1.0, 85.0]))
# 4. Optimizer
optimizer = optim.SGD(net.parameters(),
lr=config['lr'],
momentum=config['momentum'],
weight_decay=config['weight_decay'])
scheduler = lr_scheduler.StepLR(optimizer, step_size=1000, gamma=0.8)
# 5. Tensorboard logger
logger_train = Logger(logdir=os.path.join(config['log_folder'],
'fold_{}'.format(fold_idx),
'train'),
flush_secs=2)
logger_eval = Logger(logdir=os.path.join(config['log_folder'],
'fold_{}'.format(fold_idx),
'eval'),
flush_secs=2)
# 6. Train loop
DSC_MAX, IOU1_MAX, sensitivity_MAX, specificity_MAX = -1.0, -1.0, -1.0, -1.0
for epoch in range(config['num_epoch']):
train_op(net, dataloader_train, criterion, optimizer, scheduler, epoch,
logger_train)
DSC, IOU1, sensitivity, specificity = eval_op(net, dataloader_eval,
criterion, epoch,
logger_eval)
torch.save(net.state_dict(),
os.path.join(config['save_folder'], 'UNet.newest.pkl'))
if DSC_MAX <= DSC:
DSC_MAX = DSC
torch.save(net.state_dict(),
os.path.join(config['save_folder'], 'UNet.pkl'))
if IOU1_MAX <= IOU1: IOU1_MAX = IOU1
if sensitivity_MAX <= sensitivity: sensitivity_MAX = sensitivity
if specificity_MAX <= specificity: specificity_MAX = specificity
return DSC_MAX, IOU1_MAX, sensitivity_MAX, specificity_MAX, DSC, IOU1, sensitivity, specificity
def main():
torch.backends.cudnn.benchmark = True
args = getArgs()
torch.manual_seed(args.seed)
args.cuda = torch.cuda.is_available()
if args.cuda:
device = torch.device('cuda')
else:
device = torch.device('cpu')
# horovod 初始化
hvd.init()
torch.manual_seed(args.seed)
# 打印一下训练使用的配置
if hvd.rank() == 0:
print("Training with configure: ")
for arg in vars(args):
print("{}:\t{}".format(arg, getattr(args, arg)))
if not osp.exists(args.save_model_path):
os.makedirs(args.save_model_path)
# 保存训练配置
with open(osp.join(args.save_model_path, 'train-config.json'),
'w') as f:
json.dump(args.__dict__, f, indent=4)
# 设置随机种子，保证每个 GPU 上的权重初始化都一样
if args.cuda:
# Pin GPU to local rank
torch.cuda.set_device(hvd.local_rank())
# 这一句似乎没有用的吧。不过按照 horovod 的回复来说，还是加上好了。
torch.cuda.manual_seed(args.seed)
# data
dataset_train = SpineDataset(root=args.data, transform=my_transform)
# 分布式训练需要使用这个 sampler
sampler_train = DistributedSampler(dataset_train,
num_replicas=hvd.size(),
rank=hvd.rank())
dataloader_train = DataLoader(dataset_train,
batch_size=1,
sampler=sampler_train,
num_workers=args.num_workers,
pin_memory=True)
# model
if args.network == 'DeepLab':
if args.voc:
model = gcv.models.get_deeplab_resnet101_voc(pretrained=True)
elif args.ade:
model = gcv.models.get_deeplab_resnet101_ade(pretrained=True)
else:
model = gcv.models.DeepLabV3(nclass=args.num_classes,
backbone=args.backbone)
model.auxlayer.conv5[-1] = nn.Conv2d(256,
args.num_classes,
kernel_size=1)
model.head.block[-1] = nn.Conv2d(256, args.num_classes, kernel_size=1)
elif args.network == 'FCN':
if args.voc:
model = gcv.models.get_fcn_resnet101_voc(pretrained=True)
elif args.ade:
model = gcv.models.get_fcn_resnet101_ade(pretrained=True)
else:
model = gcv.models.FCN(nclass=args.num_classes,
backbone=args.backbone)
model.auxlayer.conv5[-1] = nn.Conv2d(256,
args.num_classes,
kernel_size=1)
model.head.conv5[-1] = nn.Conv2d(512, args.num_classes, kernel_size=1)
elif args.network == 'PSPNet':
if args.voc:
model = gcv.models.get_psp_resnet101_voc(pretrained=True)
elif args.ade:
model = gcv.models.get_psp_resnet101_ade(pretrained=True)
else:
model = gcv.models.PSP(nclass=args.num_classes,
backbone=args.backbone)
model.auxlayer.conv5[-1] = nn.Conv2d(256, 2, kernel_size=1)
model.head.conv5[-1] = nn.Conv2d(512, args.num_classes, kernel_size=1)
elif args.network == 'UNet':
model = UNet(n_class=args.num_classes,
backbone=args.backbone,
pretrained=True)
model = convert_syncbn_model(model)
model = model.to(device)
# optimizer 要用 hvd 的版本包一下
# optimizer = torch.optim.Adam(model.parameters(), args.learning_rate * hvd.size())
# 不同层使用不同的学习率
if args.network == 'UNet':
optimizer = torch.optim.SGD([
{
'params': model.down_blocks.parameters(),
'lr': args.learning_rate * 0.5
},
{
'params': model.bridge.parameters()
},
{
'params': model.head.parameters()
},
],
lr=args.learning_rate,
momentum=0.9,
weight_decay=0.0001)
elif args.network in ['FCN', 'PSPNet', 'DeepLab']:
optimizer = optim.SGD([{
'params': model.pretrained.parameters(),
'lr': args.learning_rate * 0.5
}, {
'params': model.auxlayer.parameters()
}, {
'params': model.head.parameters()
}],
lr=args.learning_rate,
momentum=0.9,
weight_decay=0.0001)
else:
optimizer = optim.SGD(model.parameters(),
lr=args.learning_rate,
momentum=0.9,
weight_decay=0.0001)
optimizer = hvd.DistributedOptimizer(
optimizer, named_parameters=model.named_parameters())
# 将模型和优化器的参数广播到各个 GPU 上
hvd.broadcast_parameters(model.state_dict(), root_rank=0)
hvd.broadcast_optimizer_state(optimizer, root_rank=0)
# lr scheduler
def poly_lr_scheduler(epoch, num_epochs=args.num_epochs, power=args.power):
return (1 - epoch / num_epochs)**power
lr_scheduler = LambdaLR(optimizer=optimizer, lr_lambda=poly_lr_scheduler)
def train(epoch):
model.train()
# Horovod: set epoch to sampler for shuffling.
sampler_train.set_epoch(epoch)
lr_scheduler.step()
loss_fn = nn.CrossEntropyLoss()
for batch_idx, (data, target) in enumerate(dataloader_train):
data = data.to(device).squeeze()
target = target.to(device).squeeze()
for batch_data, batch_target in zip(
torch.split(data, args.batch_size),
torch.split(target, args.batch_size)):
optimizer.zero_grad()
output = model(batch_data)
if args.network in ['FCN', 'PSPNet', 'DeepLab']:
loss = loss_fn(output[0], batch_target) \
+ 0.2*loss_fn(output[1], batch_target)
elif args.network == 'UNet':
loss = loss_fn(output, batch_target)
loss.backward()
optimizer.step()
if hvd.rank() == 0 and batch_idx % args.log_interval == 0:
print("Train loss: ", loss.item())
for epoch in range(args.num_epochs):
train(epoch)
if hvd.rank() == 0:
print("Saving model to {}".format(
osp.join(args.save_model_path,
"checkpoint-{:0>3d}.pth".format(epoch))))
torch.save({'state_dict': model.state_dict()},
osp.join(args.save_model_path,
"checkpoint-{:0>3d}.pth".format(epoch)))
help="test with specific model(default: retinanet)",
type=str,
choices=["retinanet", "unet"])
parser.add_argument(
"-v",
"--visual",
help="test with specific model and visualize the result.",
action="store_true")
parser.add_argument("-p",
"--path",
help="test with specific model weight file from path",
type=str,
default=config["RetinaNet"]["weight_load_path"])
args = parser.parse_args()
if args.model == "retinanet":
net = RetinaNet(num_classes).cuda()
net.load_state_dict(torch.load(args.path))
result = test_retinanet.test_retinanet_using_IEC(net, args.visual)
result2 = test_retinanet.test_retinanet_using_ANE_CAL(net, args.visual)
elif args.model == "unet":
net = UNet(1, 4).cuda()
net.load_state_dict(torch.load(args.path))
result = test_unet.test_unet_using_IEC(net)
else:
# default using retinanet to test
net = RetinaNet(num_classes).cuda()
net.load_state_dict(torch.load(args.path))
#result = test_retinanet.test_retinanet_using_IEC(net, args.visual)
#result2 = test_retinanet.test_retinanet_using_ANE_CAL(net, args.visual)
test_retinanet.test_retinanet_by_qrs(net)