def prepare_unet3d2(num_epochs=2, initial_lr=0.001, device=0):
device = torch.device(device)
model3d = (
UNet3D(
normalization="batch",
preactivation=True,
residual=True,
num_encoding_blocks=3,
upsampling_type="trilinear",
)
.to(device)
.eval()
)
# optimizer = optim.Adam(model3d.parameters(), lr=initial_lr)
# scheduler = optim.lr_scheduler.ExponentialLR(optimizer, num_epochs)
optimizer = optim.AdamW(
model3d.parameters(),
lr=0.001,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=0.01,
amsgrad=False,
)
scheduler = lr_scheduler.StepLR(optimizer, step_size=40, gamma=0.51)
return model3d, optimizer, scheduler
def prepare_unet3d(existing_model_fname=None, num_epochs=2, initial_lr=0.001, device=0):
from unet import UNet2D, UNet3D
device = torch.device(device)
model3d = UNet3D(
normalization="batch",
preactivation=True,
residual=True,
num_encoding_blocks=3,
upsampling_type="trilinear",
)
if existing_model_fname is not None:
model3d = load_model(existing_model_fname, model3d)
if torch.cuda.device_count() > 1:
print(f"Using {torch.cuda.device_count()} gpus.")
model3d = nn.DataParallel(model3d).to(device).eval()
else:
model3d = model3d.to(device).eval()
# optimizer = optim.Adam(model3d.parameters(), lr=initial_lr)
# scheduler = optim.lr_scheduler.ExponentialLR(optimizer, num_epochs)
optimizer = optim.AdamW(
model3d.parameters(),
lr=initial_lr,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=0.01,
amsgrad=False,
)
scheduler = lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.51)
return model3d, optimizer, scheduler
def __init__(self, hparams, tr_im, tr_lb, val_im, val_lb):
super(KitsTrainer, self).__init__()
# hparam経由では、listで渡すことは出来ない。
self.hparams = hparams
# TODO: yaml
self.tr_DS = KitsDataSet(tr_im[:100], tr_lb[:100], phase='train', transform=None)
self.val_DS = KitsDataSet(val_im[:100], val_lb[:100], phase='val', transform=None)
self.criterion = BCEDiceLoss(0.5, 0.5)
self.batch_size = hparams.batch_size
self.net = UNet3D([self.batch_size, 1, 16, 48, 48], 3).cuda()
def test_unet_3d():
model = UNet3D(
normalization='batch',
preactivation=True,
residual=False,
).to(device).eval()
shape = 1, 1, 132, 132, 116
result = 1, 2, 44, 44, 28
y = run(model, shape)
assert tuple(y.shape) == result
def test_unet_3d_residual():
model = UNet3D(
normalization='batch',
preactivation=True,
residual=True,
num_encoding_blocks=2,
upsampling_type='trilinear',
).to(device).eval()
shape = 1, 1, 64, 64, 56
result = 1, 2, 64, 64, 56
y = run(model, shape)
assert tuple(y.shape) == result
def build_model(self):
print('==> Build model and setup loss and optimizer')
#build model
self.model = UNet3D(in_channels=3, out_channels=2)
#print self.model
if self.device.type == "cuda":
self.model = torch.nn.DataParallel(self.model)
self.model.to(self.device)
#Loss function and optimizer
# self.criterion = nn.binary_cross_entropy_with_logits().cuda()
self.optimizer = torch.optim.SGD(self.model.parameters(),
self.lr,
momentum=0.9)
self.scheduler = ReduceLROnPlateau(self.optimizer,
'min',
patience=1,
verbose=True)
def get_network(num_channels, num_classes):
"""
Return the deep neural network architecture.
:param num_channels: number of input channels.
:return: neural network model.
"""
# Typical Unet architecture used in SoA pipeline
network = UNet3D(
in_channels=num_channels,
out_classes=num_classes,
out_channels_first_layer=30,
residual=False,
normalization='instance',
padding=True,
activation='LeakyReLU',
upsampling_type='trilinear',
)
# Set the network in evaluation mode
network.eval()
# Put the model on the gpu.
if torch.cuda.is_available():
network.cuda()
return network
import numpy as np
import torch
from unet import UNet3D
with torch.no_grad():
state = torch.load('unet3d-lateral-root-nuclei-lightsheet-ds1x.pytorch')
net = UNet3D(in_channels=1, out_channels=2, f_maps=32, testing=True)
net.load_state_dict(state)
# load and normalize the input
im = np.load('test_input.npz')
im = im['arr_0'].astype('float32')
im -= im.mean()
im /= im.std()
# forward pass
inp = torch.from_numpy(im)
out = net(inp)
out = out.cpu().numpy()
# compare with test_output
test_out = np.load('test_output.npz')
test_out = test_out['arr_0']
assert np.allclose(out, test_out)
'label': [
{'name': 'RandomFlip'},
{'name': 'RandomRotate90'},
{'name': 'RandomRotate', 'axes': [[2, 1]], 'angle_spectrum': 15, 'mode':'reflect'},
{'name': 'ElasticDeformation', 'spline_order': 0},
{'name': 'ToTensor', 'expand_dims': True}
]
}
with torch.no_grad():
state = torch.load('PreTrainedModels/unet3d-arabidopsis-ovules-confocal-ds2x.pytorch')
net = UNet3D(in_channels=1,
out_channels=1,
f_maps=[32, 64, 128, 256],
testing=True)
net.load_state_dict(state)
# load and normalize the input
inputDir = '/media/pablo/d7c61090-024c-469a-930c-f5ada47fb049/PabloVicenteMunuera/CellShape3DAnalysis/Datasets/HDF5/RobTetley/';
nameFile = 'Part2_Decon_c1_t1_predictions_best.h5'
with h5py.File(inputDir + nameFile) as trainingInput:
min_value, max_value, mean, std = calculate_stats(trainingInput)
transformer = transforms.get_transformer(transformer_config, min_value=min_value, max_value=max_value,
mean=mean, std=std)
# load raw images transformer
raw_transform = transformer.raw_transform()
# forward pass