def train(network_specs,
training_params,
image_path,
save_path,
ckpt_path,
epoch=10):
print('creating datapipe...')
# create images DataPipeline
datapipe = DataPipeline(image_path=image_path,
training_params=training_params)
print('creating network model...')
# create model VAE
model = UNet(network_specs=network_specs,
datapipe=datapipe,
training_params=training_params)
# train the model
# save_config is flexible
print('''
=============
HERE WE GO
=============
''')
model.train(save_path=save_path, ckpt_path=ckpt_path, epoch=epoch)
def train(network_specs, training_params, image_path, save_path, ckpt_path):
print('creating datapipe...')
# create images DataPipeline
datapipe = DataPipeline(image_path=image_path,
training_params=training_params)
print('creating network model...')
# create model VAE
model = UNet(network_specs=network_specs,
datapipe=datapipe,
training_params=training_params,
mode='evaluating')
# train the model
# save_config is flexible
print('''
=============
HERE WE GO
=============
''')
images, labels, preds = model.evaluate(ckpt_path=ckpt_path)
for i in range(labels.shape[0]):
# plt.subplots(figsize=[16,12])
for j in range(3):
plt.subplot(2, 6, 2 * j + 1)
plt.imshow(labels[i, :, :, j])
plt.subplot(2, 6, 2 * j + 2)
plt.imshow(preds[i, :, :, j])
plt.subplot(2, 6, 7)
plt.imshow(np.squeeze(images[i]))
plt.show()
def train(epochs, batch_size, learning_rate):
train_loader = torch.utils.data.DataLoader(SegThorDataset(
"data",
phase='train',
transform=transforms.Compose([Rescale(0.25),
Normalize(),
ToTensor()]),
target_transform=transforms.Compose([Rescale(0.25),
ToTensor()])),
batch_size=batch_size,
shuffle=True)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = UNet().to(device)
model.apply(weight_init)
#optimizer = optim.Adam(model.parameters(), lr=learning_rate) #learning rate to 0.001 for initial stage
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.95)
#optimizer = adabound.AdaBound(params = model.parameters(), lr = 0.001, final_lr = 0.1)
for epoch in range(epochs):
print('Epoch {}/{}'.format(epoch + 1, epochs))
print('-' * 10)
running_loss = 0.0
loss_seg = np.zeros(5)
for batch_idx, (train_data, labels) in enumerate(train_loader):
train_data, labels = train_data.to(
device, dtype=torch.float), labels.to(device,
dtype=torch.uint8)
print("train data size", train_data.size())
print("label size", labels.size())
optimizer.zero_grad()
output = model(train_data)
print("output: {} and taget: {}".format(output.size(),
labels.size()))
loss_label, loss = dice_loss(output, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
for i in range(4):
loss_seg[i] += loss_label[i]
print("Length: ", len(train_loader))
epoch_loss = running_loss / len(train_loader)
epoch_loss_class = np.true_divide(loss_seg, len(train_loader))
print(
"Dice per class: Background = {:.4f} Eusophagus = {:.4f} Heart = {:.4f} Trachea = {:.4f} Aorta = {:.4f}\n"
.format(epoch_loss_class[0], epoch_loss_class[1],
epoch_loss_class[2], epoch_loss_class[3],
epoch_loss_class[4]))
print("Total Dice Loss: {:.4f}\n".format(epoch_loss))
os.makedirs("models", exist_ok=True)
torch.save(model, "models/model.pt")
def main():
FLAGS = parser.parse_args()
# Calculate the predictions for all the images in the input_img_dir.
for img_name in os.listdir(FLAGS.input_img_dir):
if img_name.endswith('.png'):
original_img = load_img(
FLAGS.input_img_dir + '/' +
img_name) # load_img function exists in file utils.py
# Resizing image because of the small memory size
input_img = tl.prepro.imresize(original_img,
[FLAGS.img_size, FLAGS.img_size])
input_img = np.reshape(input_img,
[1, FLAGS.img_size, FLAGS.img_size, 3])
unet = UNet(FLAGS.img_size)
# The output is an array of the size(img_size * img_size, 1)
prediction = unet.predict(input_img, FLAGS.model_save_dir)
# Saving the image given the probabilities
# save_img fnuction exists in file utils.py
save_img(
prediction, original_img, FLAGS.img_size, FLAGS.input_img_dir +
'/' + img_name.split('.')[0] + '_pred.png')
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():
train_root_dir = '/content/drive/My Drive/DDSM/train/CBIS-DDSM'
test_root_dir = '/content/drive/My Drive/DDSM/test/CBIS-DDSM'
path_weights = '/content/drive/My Drive/Cv/weights'
batch_size = 3
valid_size = 0.2
nb_epochs = 20
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# data loaders
loaders = dataloaders(train_root_dir, combined_transform, batch_size, valid_size)
model = UNet(in_channels=3, out_channels=1)
model.to(device)
optimizer = optim.Adam(model.parameters(), lr=0.01)
exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=4, gamma=0.3)
model = train(model, optimizer, exp_lr_scheduler, loaders, nb_epochs, device, path_weights)
# from torchsummary import summary
#
# summary(model, input_size=(3, 224, 224))
# test_transform = transforms.Compose([
# transforms.ToTensor(),
# transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
# ])
test_loader = DataLoader(
MassSegmentationDataset(test_root_dir, combined_transform),
batch_size=batch_size,
num_workers=0
)
test(model, test_loader, device)
def main():
FLAGS = parser.parse_args()
# Loading train and test data.
# load_data function exists in file utils.py
print("Loading dataset.")
X_train, y_train = load_data(FLAGS.dataset_dir + '/train', FLAGS.img_size,
FLAGS.augment_data)
X_test, y_test = load_data(FLAGS.dataset_dir + '/test', FLAGS.img_size,
FLAGS.augment_data)
# Making sure that the data was loaded successfully.
print("Train set image size : ", X_train.shape)
print("Train set label size : ", y_train.shape)
print("Test set image size : ", X_test.shape)
print("Test set label size : ", y_test.shape)
print("Dataset loaded successfully.")
# Creating a unet object.
# class UNet exists in file model.py
unet = UNet(FLAGS.img_size)
# Training the network, printing the loss value for every epoch
# , and the accuracy on the test set after the training is complete
train_loss_values, test_loss_values = unet.train(X_train, y_train, X_test,
y_test, FLAGS.num_epochs,
FLAGS.learning_rate,
FLAGS.model_save_dir)
# Plotting loss values on train and test set, and saving it as an image Loss.png
# plot_loss exists in file utils.py
plot_loss(train_loss_values, test_loss_values, 'Loss.png')
def main(args):
train_dataloader, test_dataloader = dataloader.load_datasets(
batch_size=args.batch_size,
image_resize=args.image_resize,
train_dataset_size=args.train_data_size,
test_dataset_size=args.test_data_size,
download=args.download_dataset
)
model = UNet(out_channels=21)
optimizer = optim.Adam(model.parameters(), lr=args.lr)
ce_weight = utils.get_weight(train_dataloader.dataset)
if len(ce_weight) < 21:
criterion = nn.CrossEntropyLoss()
else:
criterion = nn.CrossEntropyLoss(utils.get_weight(train_dataloader.dataset))
print(f'Start training for {args.epochs} epochs')
train(model=model,
dataloader=train_dataloader,
epochs=args.epochs,
optimizer=optimizer,
criterion=criterion,
save_output_every=1,
)
print(f'Training finished')
print(f'Start evaluating with {len(test_dataloader.dataset)} images')
eval(model, test_dataloader)
print('All done')
def __init__(self, opt):
self.opt = opt
if opt.inference:
self.testset = TestImageDataset(fdir=opt.impaths['test'],
imsize=opt.imsize)
else:
self.trainset = ImageDataset(fdir=opt.impaths['train'],
bdir=opt.impaths['btrain'],
imsize=opt.imsize,
mode='train',
aug_prob=opt.aug_prob,
prefetch=opt.prefetch)
self.valset = ImageDataset(fdir=opt.impaths['val'],
bdir=opt.impaths['bval'],
imsize=opt.imsize,
mode='val',
aug_prob=opt.aug_prob,
prefetch=opt.prefetch)
self.model = UNet(n_channels=3,
n_classes=1,
bilinear=self.opt.use_bilinear)
if opt.checkpoint:
self.model.load_state_dict(
torch.load('./state_dict/{:s}'.format(opt.checkpoint),
map_location=self.opt.device))
print('checkpoint {:s} has been loaded'.format(opt.checkpoint))
if opt.multi_gpu == 'on':
self.model = torch.nn.DataParallel(self.model)
self.model = self.model.to(opt.device)
self._print_args()
def main():
args = parser.parse_args()
dataset = SyntheticCellDataset(arg.img_dir, arg.mask_dir)
indices = torch.randperm(len(dataset)).tolist()
sr = int(args.split_ratio * len(dataset))
train_set = torch.utils.data.Subset(dataset, indices[:-sr])
val_set = torch.utils.data.Subset(dataset, indices[-sr:])
train_loader = torch.utils.data.DataLoader(train_set, batch_size=args.batch_size, shuffle=True, pin_memory=True)
val_loader = torch.utils.data.DataLoader(val_set, batch_size=args.batch_size, shuffle=False, pin_memory=True)
device = torch.device("cpu" if not args.use_cuda else "cuda:0")
model = UNet()
model.to(device)
dsc_loss = DiceLoss()
optimizer = torch.optim.Adam(model.parameters(), args.lr)
val_overall = 1000
for epoch in args.N_epoch:
model, train_loss, optimizer = train(model, train_loader, device, optimizer)
val_loss = validate(model, val_loader, device)
if val_loss < val_overall:
save_checkpoint(args.model_save_dir + '/epoch_'+str(epoch+1), model, train_loss, val_loss, epoch)
val_overall = val_loss
print('[{}/{}] train loss :{} val loss : {}'.format(epoch+1, num_epoch, train_loss, val_loss))
print('Training completed)
def lr_find(model: UNet,
data_loader,
optimizer: Optimizer,
criterion,
use_gpu,
min_lr=0.0001,
max_lr=0.1):
# Save model and optimizer states to revert
model_state = model.state_dict()
optimizer_state = optimizer.state_dict()
losses = []
lrs = []
scheduler = CyclicExpLR(optimizer,
min_lr,
max_lr,
step_size_up=100,
mode='triangular',
cycle_momentum=True)
model.train()
for i, (data, target, class_ids) in enumerate(data_loader):
data, target = data, target
if use_gpu:
data = data.cuda()
target = target.cuda()
optimizer.zero_grad()
output_raw = model(data)
# This step is specific for this project
output = torch.zeros(output_raw.shape[0], 1, output_raw.shape[2],
output_raw.shape[3])
if use_gpu:
output = output.cuda()
# This step is specific for this project
for idx, (raw_o, class_id) in enumerate(zip(output_raw, class_ids)):
output[idx] = raw_o[class_id - 1]
loss = criterion(output, target)
loss.backward()
current_lr = optimizer.param_groups[0]['lr']
# Stop if lr stopped increasing
if len(lrs) > 0 and current_lr < lrs[-1]:
break
lrs.append(current_lr)
losses.append(loss.item())
optimizer.step()
scheduler.step()
# Plot in log scale
plt.plot(lrs, losses)
plt.xscale('log')
plt.show()
model.load_state_dict(model_state)
optimizer.load_state_dict(optimizer_state)
def build_model(self):
"""Create a model"""
self.model = UNet(self.in_dim, self.out_dim, self.num_filters)
self.model = self.model.float()
self.optimizer = torch.optim.Adam(self.model.parameters(),
self.lr, [self.beta1, self.beta2],
weight_decay=self.weight_decay)
self.print_network(self.model, 'unet')
self.model.to(self.device)
def main():
"""
Training.
"""
global start_epoch, epoch, checkpoint
# Initialize model or load checkpoint
if checkpoint is None:
model = UNet(in_channels, out_channels)
# Initialize the optimizer
optimizer = torch.optim.Adam(params=filter(lambda p: p.requires_grad,
model.parameters()),
lr=lr)
else:
checkpoint = torch.load(checkpoint)
start_epoch = checkpoint['epoch'] + 1
model = checkpoint['model']
optimizer = checkpoint['optimizer']
# Move to default device
model = model.to(device)
criterion = nn.L1Loss().to(device)
# Custom dataloaders
train_loader = torch.utils.data.DataLoader(TripletDataset(
train_folder, crop_size, scale),
batch_size=batch_size,
shuffle=True,
num_workers=workers,
pin_memory=True)
test_loader = torch.utils.data.DataLoader(TripletDataset(
test_folder, crop_size, scale),
batch_size=batch_size,
shuffle=True,
num_workers=workers,
pin_memory=True)
# Total number of epochs to train for
epochs = int(iterations // len(train_loader) + 1)
# Epochs
for epoch in range(start_epoch, epochs):
# One epoch's training
train(train_loader=train_loader,
model=model,
criterion=criterion,
optimizer=optimizer,
epoch=epoch,
epochs=epochs)
test(test_loader=test_loader, model=model, criterion=criterion)
# Save checkpoint
torch.save({
'epoch': epoch,
'model': model,
'optimizer': optimizer
}, f'checkpoints/checkpoint_unet_{epoch}.pth.tar')
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--action",
type=str,
default='train',
help="train or test")
args = parser.parse_args()
config = load_config()
# 使用tensorboard
time_now = datetime.now().isoformat()
if not os.path.exists(config.RUN_PATH):
os.mkdir(config.RUN_PATH)
writer = SummaryWriter(log_dir=config.RUN_PATH)
# 随机数种子
torch.manual_seed(config.SEED)
torch.cuda.manual_seed(config.SEED)
np.random.seed(config.SEED)
random.seed(config.SEED)
# INIT GPU
os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(str(e) for e in config.GPU)
if torch.cuda.is_available():
config.DEVICE = torch.device("cuda")
print('\nGPU IS AVAILABLE')
torch.backends.cudnn.benchmark = True
else:
config.DEVICE = torch.device("cpu")
net = UNet(2).to(config.DEVICE)
print(list(torchvision.models.resnet18(False).children())[7])
optimizer = optim.Adam(net.parameters(), betas=(0.5, 0.999), lr=config.LR)
loss = nn.L1Loss()
# 加载数据集
if args.action == 'train':
train_dataset = LABDataset(config, config.TRAIN_PATH)
len_train = len(train_dataset)
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=config.BATCH_SIZE, shuffle=True)
iter_per_epoch = len(train_loader)
train_(config, train_loader, net, optimizer, loss, len_train,
iter_per_epoch, writer)
if args.action == "test":
test_dataset = LABDataset(config, config.TEST_PATH)
test_loader = torch.utils.data.DataLoader(test_dataset,
batch_size=1,
shuffle=False)
test(config, test_loader, net, loss)
def __init__(self,
cur_dir,
suffix='.tif',
cuda=True,
testBatchSize=4,
batchSize=4,
nEpochs=200,
lr=0.01,
threads=4,
seed=123,
size=256,
input_transform=True,
target_transform=True):
# super(TrainModel, self).__init__()
self.data_dir = cur_dir + '/data/'
self.suffix = suffix
"""
training parameters are set here
"""
self.colordim = 1
self.cuda = cuda
if self.cuda and not torch.cuda.is_available():
raise Exception("No GPU found, please run without --cuda")
self.testBatchSize = testBatchSize
self.batchSize = batchSize
self.nEpochs = nEpochs
self.lr = lr
self.threads = threads
self.seed = seed
self.size = size
self.input_transform = input_transform
self.target_transform = target_transform
self.__check_dir = cur_dir + '/checkpoint'
if not exists(self.__check_dir):
os.mkdir(self.__check_dir)
self.__epoch_dir = cur_dir + '/epoch'
if not exists(self.__epoch_dir):
os.mkdir(self.__epoch_dir)
"""
initialize the model
"""
if self.cuda:
self.unet = UNet(self.colordim).cuda()
self.criterion = nn.MSELoss().cuda()
else:
self.unet = UNet(self.colordim)
self.criterion = nn.MSELoss()
self.optimizer = optim.SGD(self.unet.parameters(),
lr=self.lr,
momentum=0.9,
weight_decay=0.0001)
def main():
train_transform = A.Compose([
A.Resize(height=config.IMAGE_HEIGHT, width=config.IMAGE_WIDTH),
A.Rotate(limit=35, p=1.0),
A.HorizontalFlip(p=0.5),
A.VerticalFlip(p=0.1),
A.normalize(mean=[0.0, 0.0, 0.0],
std=[1.0, 1.0, 1.0],
max_pixel_value=255.0),
ToTensorV2,
])
val_transform = A.Compose([
A.Resize(height=config.IMAGE_HEIGHT, width=config.IMAGE_WIDTH),
A.normalize(mean=[0.0, 0.0, 0.0],
std=[1.0, 1.0, 1.0],
max_pixel_value=255.0),
ToTensorV2,
])
model = UNet(in_channels=3, out_channels=1).to(config.DEVICE)
loss_fn = nn.BCEWithLogitsLoss()
optimizer = optim.Adam(model.parameters(), lr=config.LEARNING_RATE)
train_loader, val_loader = get_loaders(
config.TRAIN_IMAGE_DIR,
config.TRAIN_MASK_DIR,
config.VAL_IMG_DIR,
config.VAL_MASK_DIR,
config.BATCH_SIZE,
train_transform,
val_transform,
)
if config.LOAD_MODEL:
load_checkpoint(torch.load('my_checkpoint.pth.tar'), model)
scaler = torch.cuda.amp.GradScaler()
for epoch in range(config.NUM_EPOCHS):
train_fn(train_loader, model, optimizer, loss_fn, scaler)
# save model
checkpoint = {
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
}
save_checkpoint(checkpoint)
# check acc
check_accuracy(val_loader, model, device=config.DEVICE)
# print some examples to a folder
save_predictions_as_imgs(val_loader,
model,
folder='saved_images',
device=config.DEVICE)
def main():
# width_in = 284
# height_in = 284
# width_out = 196
# height_out = 196
# PATH = './unet.pt'
# x_train, y_train, x_val, y_val = get_dataset(width_in, height_in, width_out, height_out)
# print(x_train.shape, y_train.shape, x_val.shape, y_val.shape)
batch_size = 3
epochs = 1
epoch_lapse = 50
threshold = 0.5
learning_rate = 0.01
unet = UNet(in_channel=1, out_channel=2)
if use_gpu:
unet = unet.cuda()
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(unet.parameters(), lr=0.01, momentum=0.99)
if sys.argv[1] == 'train':
train(unet, batch_size, epochs, epoch_lapse, threshold, learning_rate,
criterion, optimizer, x_train, y_train, x_val, y_val, width_out,
height_out)
pass
else:
if use_gpu:
unet.load_state_dict(torch.load(PATH))
else:
unet.load_state_dict(torch.load(PATH, map_location='cpu'))
print(unet.eval())
def create_net(self) -> nn.Module:
if self.model_architecture == 'unet':
self.net = UNet(n_channels=1, n_classes=1)
elif self.model_architecture == 'mnet2':
self.net = MobileNetV2_UNet()
else:
raise ValueError(
f'model_architecture must be in ["unet", "mnet2"]. '
f'passed: {self.model_architecture}')
self.net.to(device=self.device)
return self.net
def __init__(self,hparams):
super(Unetmodel,self).__init__()
self.hparams = hparams
self.hdf_path = self.hparams.hdf_path
self.max_dist = self.hparams.max_dist
self.grid_resolution = self.hparams.grid_resolution
self.augment = self.hparams.augment
self.net = UNet()
self.metric = Ovl('Volume_overlap')
self.loss = Dice_loss()
self.lr = self.hparams.lr
self.batch_size = self.hparams.batch_size
def __init__(self, device, n_channels=1, n_classes=1, load_model=None):
self.device = device
self.name = "unet"
self.mask_type = torch.float32 if n_classes == 1 else torch.long
self.net = UNet(n_channels=n_channels, n_classes=n_classes)
if load_model is not None:
self.net.load_state_dict(torch.load(load_model,
map_location=device),
strict=False)
self.net = self.net.to(device)
self.criterion = nn.BCELoss(
) if self.net.n_classes == 1 else nn.CrossEntropyLoss()
def main(mode):
config = load_config()
# 使用tensorboard
time_now = datetime.now().isoformat()
if not os.path.exists(config.RUN_PATH):
os.mkdir(config.RUN_PATH)
writer = SummaryWriter(log_dir=config.RUN_PATH)
# 随机数种子
torch.manual_seed(config.SEED)
torch.cuda.manual_seed(config.SEED)
np.random.seed(config.SEED)
random.seed(config.SEED)
# INIT GPU
os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(str(e) for e in config.GPU)
if torch.cuda.is_available():
config.DEVICE = torch.device("cuda")
print('\nGPU IS AVAILABLE')
torch.backends.cudnn.benchmark = True
else:
config.DEVICE = torch.device("cpu")
net = UNet(1).to(config.DEVICE)
optimizer = optim.Adam(net.parameters(), betas=(0.5, 0.999), lr=config.LR)
#criterion = nn.CrossEntropyLoss() # 定义loss函数
criterion = nn.BCELoss()
# 加载数据集
if mode == 1:
train_dataset = MyDataset(config, config.TRAIN_PATH)
len_train = len(train_dataset)
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=config.BATCH_SIZE, shuffle=True)
iter_per_epoch = len(train_loader)
train_(config, train_loader, net, optimizer, criterion, len_train,
iter_per_epoch, writer)
if mode == 2:
test_dataset = MyDataset(config, config.TEST_PATH)
test_loader = torch.utils.data.DataLoader(test_dataset,
batch_size=1,
shuffle=False)
test(config, test_loader, net, criterion)
def build_model(self):
""" Rough """
if self.cfg.task == 'cls':
self.model = BinaryClassifier(num_classes=2)
elif self.cfg.task == 'seg':
self.model = UNet(num_classes=2)
self.criterion = nn.CrossEntropyLoss()
self.optim = torch.optim.Adam(self.model.parameters(),
lr=self.cfg.lr,
betas=(self.cfg.beta0, self.cfg.beta1))
if self.n_gpus > 1:
print('### {} of gpus are used!!!'.format(self.n_gpus))
self.model = nn.DataParallel(self.model)
self.model = self.model.to(self.device)
def get_score():
ckpt_dir = os.listdir(
'/media/muyun99/DownloadResource/dataset/opends-Supervisely Person Dataset/checkpoints'
)
Unet_dir = [
"25_noise", "50_noise", "75_noise", "25_noise_pro", "50_noise_pro",
"75_noise_pro", "75_noise_pro_finetune"
]
UNet_Pick_dir = ["75_noise_pro_QAM", "75_noise_pro_QAM_finetune"]
UNet_Pick_cbam_dir = ["75_noise_pro_QAM_cbam_finetune"]
for dir in ckpt_dir:
print(dir)
flag = "unet_pick"
if dir in Unet_dir:
net = UNet(n_classes=1, n_channels=3)
flag = "unet"
elif dir in UNet_Pick_dir:
net = UNet_Pick(n_classes=1, n_channels=3)
elif dir in UNet_Pick_cbam_dir:
net = UNet_Pick_cbam(n_classes=1, n_channels=3)
else:
continue
true_dir = os.path.join(
'/media/muyun99/DownloadResource/dataset/opends-Supervisely Person Dataset/checkpoints',
dir)
try:
best_ckpt_path = get_best_checkpoint(net, true_dir, flag)
test_score = get_test_score(net, best_ckpt_path, flag)
print(f'{dir} best test score is {test_score}')
except Exception as ex:
print(f'{dir} error')
print(f'出现异常 {ex}')
continue
def train(args):
model = UNet(n_channels=125, n_classes=10).to(device)
batch_size = args.batch_size
num_epochs = args.num_epochs
save_path = args.save_path
criterion = L.CrossEntropyLoss2d() # 损失函数
optimizer = optim.Adam(model.parameters(
)) # 梯度下降 # model.parameters():Returns an iterator over module parameters
liver_dataset = LiverDataset("data",
transform=x_transform,
target_transform=y_transform) # 加载数据集
dataloader = DataLoader(liver_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=8)
train_model(model, criterion, optimizer, dataloader, num_epochs, save_path)
test(args)
def __init__(self,config,trainLoader,validLoader):
self.config = config
self.trainLoader = trainLoader
self.validLoader = validLoader
self.numTrain = len(self.trainLoader.dataset)
self.numValid = len(self.validLoader.dataset)
self.saveModelDir = str(self.config.save_model_dir)+"/"
self.bestModel = config.bestModel
self.useGpu = self.config.use_gpu
self.net = UNet()
if(self.config.resume == True):
print("LOADING SAVED MODEL")
self.loadCheckpoint()
else:
print("INTIALIZING NEW MODEL")
self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
self.net = self.net.to(self.device)
self.totalEpochs = config.epochs
self.optimizer = optim.Adam(self.net.parameters(), lr=5e-4)
self.loss = DiceLoss()
self.num_params = sum([p.data.nelement() for p in self.net.parameters()])
self.trainPaitence = config.train_paitence
if not self.config.resume: # self.freezeLayers(6)
summary(self.net, input_size=(3,256,256))
print('[*] Number of model parameters: {:,}'.format(self.num_params))
self.writer = SummaryWriter(self.config.tensorboard_path+"/")
def __init__(self, n_channels, n_classes):
super(UNet_Pick_cbam, self).__init__()
self.n_channels = n_channels
self.n_classes = n_classes
self.Unet = UNet(n_classes=n_classes, n_channels=n_channels)
self.QAM = implement.QAM_cbam()
self.OCM = implement.OCM()
def main(conf):
device = "cuda:0" if torch.cuda.is_available() else 'cpu'
beta_schedule = "linear"
beta_start = 1e-4
beta_end = 2e-2
n_timestep = 1000
conf.distributed = dist.get_world_size() > 1
transform = transforms.Compose(
[
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True),
]
)
train_set = MultiResolutionDataset(
conf.dataset.path, transform, conf.dataset.resolution
)
train_sampler = dist.data_sampler(
train_set, shuffle=True, distributed=conf.distributed
)
train_loader = conf.training.dataloader.make(train_set, sampler=train_sampler)
model = UNet(
conf.model.in_channel,
conf.model.channel,
channel_multiplier=conf.model.channel_multiplier,
n_res_blocks=conf.model.n_res_blocks,
attn_strides=conf.model.attn_strides,
dropout=conf.model.dropout,
fold=conf.model.fold,
)
model = model.to(device)
ema = UNet(
conf.model.in_channel,
conf.model.channel,
channel_multiplier=conf.model.channel_multiplier,
n_res_blocks=conf.model.n_res_blocks,
attn_strides=conf.model.attn_strides,
dropout=conf.model.dropout,
fold=conf.model.fold,
)
ema = ema.to(device)
if conf.distributed:
model = nn.parallel.DistributedDataParallel(
model,
device_ids=[dist.get_local_rank()],
output_device=dist.get_local_rank(),
)
optimizer = conf.training.optimizer.make(model.parameters())
scheduler = conf.training.scheduler.make(optimizer)
betas = make_beta_schedule(beta_schedule, beta_start, beta_end, n_timestep)
diffusion = GaussianDiffusion(betas).to(device)
train(conf, train_loader, model, ema, diffusion, optimizer, scheduler, device)
def load_model(data, model_path, cuda=True):
if cuda and not torch.cuda.is_available():
raise Exception("No GPU found, please run without --cuda")
unet = UNet()
if cuda:
unet = unet.cuda()
if not cuda:
unet.load_state_dict(
torch.load(model_path, map_location=lambda storage, loc: storage))
else:
unet.load_state_dict(torch.load(model_path))
if cuda:
data = Variable(data.cuda())
else:
data = Variable(data)
data = torch.unsqueeze(data, 0)
output = unet(data)
if cuda:
output = output.cuda()
return output
def main():
args = parser.parse_args()
device = torch.device("cpu" if not args.use_cuda else "cuda:0")
model = UNet()
load_checkpoint(args.weight, model, device='cpu')
model.to(device)
img = Image.open(img_file)
resize = transforms.Resize(size=(576, 576))
im_r = TF.to_tensor(resize(img))
im_r = im_r.unsqueeze(0)
with torch.no_grad():
pred = model(im_r.to(device))
pred_mask = pred.detach().cpu().numpy().squeeze()
def main(tocsv=False, save=False, mask=False, valid_train=False, toiou=False):
model_name = config['param']['model']
resize = not config['valid'].getboolean('pred_orig_size')
if model_name == 'unet_vgg16':
model = UNetVgg16(3, 1, fixed_vgg=True)
elif model_name == 'dcan':
model = DCAN(3, 1)
elif model_name == 'caunet':
model = CAUNet()
elif model_name == 'camunet':
model = CAMUNet()
else:
model = UNet()
if torch.cuda.is_available():
model = model.cuda()
# model = torch.nn.DataParallel(model).cuda()
# Sets the model in evaluation mode.
model.eval()
epoch = load_ckpt(model)
if epoch == 0:
print("Aborted: checkpoint not found!")
return
# prepare dataset
compose = Compose(augment=False, resize=resize)
data_dir = 'data/stage1_train' if valid_train else 'data/stage1_test'
dataset = KaggleDataset(data_dir, transform=compose)
iter = predict(model, dataset, compose, resize)
if tocsv:
with open('result.csv', 'w') as csvfile:
writer = csv.writer(csvfile)
writer.writerow(['ImageId', 'EncodedPixels'])
for uid, _, y, y_c, y_m, _, _, _, _ in iter:
for rle in prob_to_rles(y, y_c, y_m):
writer.writerow([uid, ' '.join([str(i) for i in rle])])
elif toiou and valid_train:
with open('iou.csv', 'w') as csvfile:
writer = csv.writer(csvfile)
writer.writerow(['ImageId', 'IoU'])
for uid, _, y, y_c, y_m, gt, _, _, _ in tqdm(iter):
iou = get_iou(y, y_c, y_m, gt)
writer.writerow([uid, iou])
else:
for uid, x, y, y_c, y_m, gt, gt_s, gt_c, gt_m in tqdm(iter):
if valid_train:
show_groundtruth(uid, x, y, y_c, y_m, gt, gt_s, gt_c, gt_m,
save)
elif mask:
save_mask(uid, y, y_c, y_m)
else:
show(uid, x, y, y_c, y_m, save)
