def __init__(self, num_classes, backbone_fn, chip_size=224):
super().__init__()
if getattr(backbone_fn, '_is_multispectral', False):
self.backbone = create_body(backbone_fn,
pretrained=True,
cut=_get_backbone_meta(
backbone_fn.__name__)['cut'])
else:
self.backbone = create_body(backbone_fn, pretrained=True)
backbone_name = backbone_fn.__name__
## Support for different backbones
if "densenet" in backbone_name or "vgg" in backbone_name:
hookable_modules = list(self.backbone.children())[0]
else:
hookable_modules = list(self.backbone.children())
if "vgg" in backbone_name:
modify_dilation_index = -5
else:
modify_dilation_index = -2
if backbone_name == 'resnet18' or backbone_name == 'resnet34':
module_to_check = 'conv'
else:
module_to_check = 'conv2'
## Hook at the index where we need to get the auxillary logits out
self.hook = hook_output(hookable_modules[modify_dilation_index])
custom_idx = 0
for i, module in enumerate(hookable_modules[modify_dilation_index:]):
dilation = 2 * (i + 1)
padding = 2 * (i + 1)
for n, m in module.named_modules():
if module_to_check in n:
m.dilation, m.padding, m.stride = (dilation, dilation), (
padding, padding), (1, 1)
elif 'downsample.0' in n:
m.stride = (1, 1)
if "vgg" in backbone_fn.__name__:
if isinstance(module, nn.Conv2d):
dilation = 2 * (custom_idx + 1)
padding = 2 * (custom_idx + 1)
module.dilation, module.padding, module.stride = (
dilation, dilation), (padding, padding), (1, 1)
custom_idx += 1
## returns the size of various activations
feature_sizes = model_sizes(self.backbone, size=(chip_size, chip_size))
## Geting the number of channel persent in stored activation inside of the hook
num_channels_aux_classifier = self.hook.stored.shape[1]
## Get number of channels in the last layer
num_channels_classifier = feature_sizes[-1][1]
self.classifier = DeepLabHead(num_channels_classifier, num_classes)
self.aux_classifier = FCNHead(num_channels_aux_classifier, num_classes)
def get_model(self, pretrained=True):
if self.model_number == 18:
body = create_body(
resnet18, pretrained=pretrained, n_in=self.input_channels, cut=-2
)
else:
body = create_body(
resnet34, pretrained=pretrained, n_in=self.input_channels, cut=-2
)
net_G = DynamicUnet(
body, self.output_channels, (self.image_size, self.image_size)
)
return net_G
def __init__(self,
base_arch,
no_diseases,
dropout=0.5,
init=nn.init.kaiming_normal_):
super(CNNPretrainedModel, self).__init__()
self.body = create_body(base_arch, pretrained=True)
nf = num_features_model(nn.Sequential(*self.body.children())) * 2
self.disease_head = create_head(nf,
no_diseases,
ps=0.5,
concat_pool=True,
bn_final=False)
#self.age_head = create_head(nf, 1, ps=0.5, concat_pool=True, bn_final=False)
self.gender_head = create_head(nf,
2,
ps=0.5,
concat_pool=True,
bn_final=False)
#self.projection_head = create_head(nf, 3, ps=0.5, concat_pool=True, bn_final=False)
self.disease_model = nn.Sequential(self.body, self.disease_head)
self.meta = cnn_config(base_arch)
self.split(self.meta['split'])
self.freeze()
apply_init(self.disease_head, init)
#apply_init(self.age_head, init)
apply_init(self.gender_head, init)
def __init__(self,
data,
scales=None,
ratios=None,
backbone=None,
pretrained_path=None,
*args,
**kwargs):
# Set default backbone to be 'resnet50'
if backbone is None:
backbone = models.resnet50
super().__init__(data, backbone)
n_bands = len(getattr(self._data, '_extract_bands', [0, 1, 2]))
_backbone = self._backbone
if hasattr(self, '_orig_backbone'):
_backbone = self._orig_backbone
# Check if a backbone provided is compatible, use resnet50 as default
if not self._check_backbone_support(_backbone):
raise Exception(
f"Enter only compatible backbones from {', '.join(self.supported_backbones)}"
)
self.name = "RetinaNet"
self._code = code
self.scales = ifnone(scales, [1, 2**(-1 / 3), 2**(-2 / 3)])
self.ratios = ifnone(ratios, [1 / 2, 1, 2])
self._n_anchors = len(self.scales) * len(self.ratios)
self._data = data
self._chip_size = (data.chip_size, data.chip_size)
# Cut-off the backbone before the penultimate layer
self._encoder = create_body(self._backbone, -2)
# Initialize the model, loss function and the Learner object
self._model = RetinaNetModel(self._encoder,
n_classes=data.c - 1,
final_bias=-4,
chip_size=self._chip_size,
n_anchors=self._n_anchors,
n_bands=n_bands)
self._loss_f = RetinaNetFocalLoss(sizes=self._model.sizes,
scales=self.scales,
ratios=self.ratios)
self.learn = Learner(data, self._model, loss_func=self._loss_f)
self.learn.split([self._model.encoder[6], self._model.c5top5])
self.learn.freeze()
if pretrained_path is not None:
self.load(str(pretrained_path))
self._arcgis_init_callback() # make first conv weights learnable
def __init__(self):
super().__init__()
concat_pool = True
self.encoder = create_body(models.resnet18, True)
nf = num_features_model(nn.Sequential(
*(self.encoder).children())) * (2 if concat_pool else 1)
self.domain_head_1, self.domain_head_2 = create_domain_head(
nf, 2, lin_ftrs=[256])
self.class_head_1, self.class_head_2 = create_classification_head(
nf, 2, lin_ftrs=[256])
self.alpha = 0
def build_res_unet(n_input=1, n_output=2, size=256):
'''
Creates a generator by building a Unet model with a Resnet18 backbone
Args:
n_input (int): Corresponds to the L feature dimension in L*a*b* color space
n_output (int): Corresponds to the ab feature dimensions in L*a*b* color space
size (int): height/width of image
Returns:
Generator model
'''
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
body = create_body(resnet18, pretrained=True, n_in=n_input, cut=-2)
net_G = DynamicUnet(body, n_output, (size, size)).to(device)
return net_G
def unet_learner_wide(
data: DataBunch,
arch: Callable,
pretrained: bool = True,
blur_final: bool = True,
norm_type: Optional[NormType] = NormType,
split_on: Optional[SplitFuncOrIdxList] = None,
blur: bool = False,
self_attention: bool = False,
y_range: Optional[Tuple[float, float]] = None,
last_cross: bool = True,
bottle: bool = False,
nf_factor: int = 1,
**kwargs: Any
) -> Learner:
"Build Unet learner from `data` and `arch`."
meta = cnn_config(arch)
body = create_body(arch, pretrained)
model = to_device(
DynamicUnetWide(
body,
n_classes=data.c,
blur=blur,
blur_final=blur_final,
self_attention=self_attention,
y_range=y_range,
norm_type=norm_type,
last_cross=last_cross,
bottle=bottle,
nf_factor=nf_factor,
),
data.device,
)
learn = Learner(data, model, **kwargs)
learn.split(ifnone(split_on, meta['split']))
if pretrained:
learn.freeze()
apply_init(model[2], nn.init.kaiming_normal_)
return learn
def unet4mri_learner(data: DataBunch,
arch: Callable,
pretrained: bool = True,
blur_final: bool = True,
norm_type: Optional[NormType] = None,
split_on: Optional[SplitFuncOrIdxList] = None,
blur: bool = False,
self_attention: bool = False,
y_range: Optional[Tuple[float, float]] = None,
last_cross: bool = True,
bottle: bool = False,
cut: Union[int, Callable] = None,
**learn_kwargs: Any) -> Learner:
"Build Unet learner from `data` and `arch`."
meta = ln.cnn_config(arch)
body = ln.create_body(arch, pretrained, cut)
try:
size = data.train_ds[0][0].size
except:
size = next(iter(data.train_dl))[0].shape[-2:]
model = to_device(
MyDynamicUnet(body,
bs=data.batch_size,
n_classes=data.c,
img_size=size,
blur=blur,
blur_final=blur_final,
self_attention=self_attention,
y_range=y_range,
norm_type=norm_type,
last_cross=last_cross,
bottle=bottle), data.device)
learn = Learner(data, model, **learn_kwargs)
learn.split(ifnone(split_on, meta['split']))
if pretrained: learn.freeze()
apply_init(model[2], nn.init.kaiming_normal_)
return learn
import argparse
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
'-i','--image', type=str, default='data/images/Test/xkObNU.jpg'
)
parser.add_argument(
'-m', '--model', default='model/color_gen.pt'
)
args = parser.parse_args()
# Load model
body = create_body(resnet18, pretrained=True, n_in=1, cut=-2)
net_G = DynamicUnet(body, 2, (256, 256)).to('cuda')
net_G.load_state_dict(torch.load(args.model))
net_G.to('cuda')
net_G.eval()
img = cv2.imread(args.image)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
lab_img = rgb2lab(img).astype(np.float32)
lab_img = T.ToTensor()(lab_img)
L = lab_img[[0], ...] / 50. - 1. # Between -1 and 1
with torch.no_grad():
pred_ab = net_G(L.unsqueeze(0).to('cuda'))
def __init__(self, data, backbone=None, pretrained_path=None):
super().__init__(data, backbone)
if self._is_multispectral:
self._backbone_ms = self._backbone
self._backbone = self._orig_backbone
scaled_mean_values = data._scaled_mean_values[
data._extract_bands].tolist()
scaled_std_values = data._scaled_std_values[
data._extract_bands].tolist()
if backbone is None:
self._backbone = models.resnet50
elif type(backbone) is str:
self._backbone = getattr(models, backbone)
else:
self._backbone = backbone
if not self._check_backbone_support(self._backbone):
raise Exception(
f"Enter only compatible backbones from {', '.join(self.supported_backbones)}"
)
self._code = instance_detector_prf
if self._backbone.__name__ is 'resnet50':
model = models.detection.maskrcnn_resnet50_fpn(
pretrained=True,
min_size=1.5 * data.chip_size,
max_size=2 * data.chip_size)
if self._is_multispectral:
model.backbone = _change_tail(model.backbone, data)
model.transform.image_mean = scaled_mean_values
model.transform.image_std = scaled_std_values
elif self._backbone.__name__ in ['resnet18', 'resnet34']:
if self._is_multispectral:
backbone_small = create_body(self._backbone_ms,
cut=_get_backbone_meta(
backbone_fn.__name__)['cut'])
backbone_small.out_channels = 512
model = models.detection.MaskRCNN(
backbone_small,
91,
min_size=1.5 * data.chip_size,
max_size=2 * data.chip_size,
image_mean=scaled_mean_values,
image_std=scaled_std_values)
else:
backbone_small = create_body(self._backbone)
backbone_small.out_channels = 512
model = models.detection.MaskRCNN(backbone_small,
91,
min_size=1.5 *
data.chip_size,
max_size=2 * data.chip_size)
else:
backbone_fpn = resnet_fpn_backbone(self._backbone.__name__, True)
if self._is_multispectral:
backbone_fpn = _change_tail(backbone_fpn, data)
model = models.detection.MaskRCNN(
backbone_fpn,
91,
min_size=1.5 * data.chip_size,
max_size=2 * data.chip_size,
image_mean=scaled_mean_values,
image_std=scaled_std_values)
else:
model = models.detection.MaskRCNN(backbone_fpn,
91,
min_size=1.5 *
data.chip_size,
max_size=2 * data.chip_size)
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, data.c)
in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
hidden_layer = 256
model.roi_heads.mask_predictor = MaskRCNNPredictor(
in_features_mask, hidden_layer, data.c)
self.learn = Learner(data, model, loss_func=mask_rcnn_loss)
self.learn.callbacks.append(train_callback(self.learn))
self.learn.model = self.learn.model.to(self._device)
self.learn.c_device = self._device
# fixes for zero division error when slice is passed
idx = 27
if self._backbone.__name__ in ['resnet18', 'resnet34']:
idx = self._freeze()
self.learn.layer_groups = split_model_idx(self.learn.model, [idx])
self.learn.create_opt(lr=3e-3)
# make first conv weights learnable
self._arcgis_init_callback()
if pretrained_path is not None:
self.load(pretrained_path)
if self._is_multispectral:
self._orig_backbone = self._backbone
self._backbone = self._backbone_ms
def __init__(self,
data,
grids=[4, 2, 1],
zooms=[0.7, 1., 1.3],
ratios=[[1., 1.], [1., 0.5], [0.5, 1.]],
backbone=None,
drop=0.3,
bias=-4.,
focal_loss=False,
pretrained_path=None,
location_loss_factor=None):
super().__init__()
self._device = torch.device(
'cuda') if torch.cuda.is_available() else torch.device('cpu')
# assert (location_loss_factor is not None) or ((location_loss_factor > 0) and (location_loss_factor < 1)),
if location_loss_factor is not None:
if not ((location_loss_factor > 0) and (location_loss_factor < 1)):
raise Exception(
'`location_loss_factor` should be greater than 0 and less than 1'
)
self.location_loss_factor = location_loss_factor
if not HAS_FASTAI:
_raise_fastai_import_error()
if backbone is None:
self._backbone = models.resnet34
elif type(backbone) is str:
self._backbone = getattr(models, backbone)
else:
self._backbone = backbone
self._create_anchors(grids, zooms, ratios)
feature_sizes = model_sizes(create_body(self._backbone),
size=(data.chip_size, data.chip_size))
num_features = feature_sizes[-1][-1]
num_channels = feature_sizes[-1][1]
ssd_head = SSDHead(grids,
self._anchors_per_cell,
data.c,
num_features=num_features,
drop=drop,
bias=bias,
num_channels=num_channels)
self._data = data
self.learn = create_cnn(data=data,
arch=self._backbone,
custom_head=ssd_head)
self.learn.model = self.learn.model.to(self._device)
if pretrained_path is not None:
self.load(pretrained_path)
if focal_loss:
self._loss_f = FocalLoss(data.c)
else:
self._loss_f = BCE_Loss(data.c)
self.learn.loss_func = self._ssd_loss
def _pspnet_unet(num_classes, backbone_fn, chip_size=224, pyramid_sizes=(1, 2, 3, 6), pretrained=True):
"""
Function which returns PPM module attached to backbone which is then used to form the Unet.
"""
if getattr(backbone_fn, '_is_multispectral', False):
backbone = create_body(backbone_fn, pretrained=pretrained, cut=_get_backbone_meta(backbone_fn.__name__)['cut'])
else:
backbone = create_body(backbone_fn, pretrained=pretrained)
backbone_name = backbone_fn.__name__
## Support for different backbones
if "densenet" in backbone_name or "vgg" in backbone_name:
hookable_modules = list(backbone.children())[0]
else:
hookable_modules = list(backbone.children())
if "vgg" in backbone_name:
modify_dilation_index = -5
else:
modify_dilation_index = -2
if backbone_name == 'resnet18' or backbone_name == 'resnet34':
module_to_check = 'conv'
else:
module_to_check = 'conv2'
custom_idx = 0
for i, module in enumerate(hookable_modules[modify_dilation_index:]):
dilation = 2 * (i + 1)
padding = 2 * (i + 1)
# padding = 1
for n, m in module.named_modules():
if module_to_check in n:
m.dilation, m.padding, m.stride = (dilation, dilation), (padding, padding), (1, 1)
elif 'downsample.0' in n:
m.stride = (1, 1)
if "vgg" in backbone_fn.__name__:
if isinstance(module, nn.Conv2d):
dilation = 2 * (custom_idx + 1)
padding = 2 * (custom_idx + 1)
module.dilation, module.padding, module.stride = (dilation, dilation), (padding, padding), (1, 1)
custom_idx += 1
## returns the size of various activations
feature_sizes = model_sizes(backbone, size=(chip_size, chip_size))
## Get number of channels in the last layer
num_channels = feature_sizes[-1][1]
penultimate_channels = num_channels / len(pyramid_sizes)
ppm = _PyramidPoolingModule(num_channels, int(penultimate_channels), pyramid_sizes)
in_final = int(penultimate_channels) * len(pyramid_sizes) + num_channels
# Reduce channel size after pyramid pooling module to avoid CUDA OOM error.
final_conv = nn.Conv2d(in_channels=in_final, out_channels=512, kernel_size=3, padding=1)
## To make Dynamic Unet work as it expects a backbone which can be indexed.
if "densenet" in backbone_name or "vgg" in backbone_name:
backbone = backbone[0]
layers = [*backbone, ppm, final_conv]
return nn.Sequential(*layers)
def __init__(self, num_classes, backbone_fn, chip_size=224, pyramid_sizes=(1, 2, 3, 6), pretrained=True):
super(PSPNet, self).__init__()
if getattr(backbone_fn, '_is_multispectral', False):
self.backbone = create_body(backbone_fn, pretrained=pretrained, cut=_get_backbone_meta(backbone_fn.__name__)['cut'])
else:
self.backbone = create_body(backbone_fn, pretrained=pretrained)
backbone_name = backbone_fn.__name__
## Support for different backbones
if "densenet" in backbone_name or "vgg" in backbone_name:
hookable_modules = list(self.backbone.children())[0]
else:
hookable_modules = list(self.backbone.children())
if "vgg" in backbone_name:
modify_dilation_index = -5
else:
modify_dilation_index = -2
if backbone_name == 'resnet18' or backbone_name == 'resnet34':
module_to_check = 'conv'
else:
module_to_check = 'conv2'
## Hook at the index where we need to get the auxillary logits out
self.hook = hook_output(hookable_modules[modify_dilation_index])
custom_idx = 0
for i, module in enumerate(hookable_modules[modify_dilation_index:]):
dilation = 2 * (i + 1)
padding = 2 * (i + 1)
for n, m in module.named_modules():
if module_to_check in n:
m.dilation, m.padding, m.stride = (dilation, dilation), (padding, padding), (1, 1)
elif 'downsample.0' in n:
m.stride = (1, 1)
if "vgg" in backbone_fn.__name__:
if isinstance(module, nn.Conv2d):
dilation = 2 * (custom_idx + 1)
padding = 2 * (custom_idx + 1)
module.dilation, module.padding, module.stride = (dilation, dilation), (padding, padding), (1, 1)
custom_idx += 1
## returns the size of various activations
feature_sizes = model_sizes(self.backbone, size=(chip_size, chip_size))
## Geting the stored parameters inside of the hook
aux_in_channels = self.hook.stored.shape[1]
## Get number of channels in the last layer
num_channels = feature_sizes[-1][1]
penultimate_channels = num_channels / len(pyramid_sizes)
self.ppm = _PyramidPoolingModule(num_channels, int(penultimate_channels), pyramid_sizes)
self.final = nn.Sequential(
## To handle case when the length of pyramid_sizes is odd
nn.Conv2d(int(penultimate_channels) * len(pyramid_sizes) + num_channels, math.ceil(penultimate_channels), kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(math.ceil(penultimate_channels)),
nn.ReLU(inplace=True),
nn.Dropout(0.1),
nn.Conv2d(math.ceil(penultimate_channels), num_classes, kernel_size=1)
)
self.aux_logits = nn.Conv2d(aux_in_channels, num_classes, kernel_size=1)
initialize_weights(self.aux_logits)
initialize_weights(self.ppm, self.final)
def __init__(self,
data,
grids=None,
zooms=[1.],
ratios=[[1., 1.]],
backbone=None,
drop=0.3,
bias=-4.,
focal_loss=False,
pretrained_path=None,
location_loss_factor=None,
ssd_version=2,
backend='pytorch'):
super().__init__(data, backbone)
self._backend = backend
if self._backend == 'tensorflow':
self._intialize_tensorflow(data, grids, zooms, ratios, backbone,
drop, bias, pretrained_path,
location_loss_factor)
else:
# assert (location_loss_factor is not None) or ((location_loss_factor > 0) and (location_loss_factor < 1)),
if not ssd_version in [1, 2]:
raise Exception("ssd_version can be only [1,2]")
if location_loss_factor is not None:
if not ((location_loss_factor > 0) and
(location_loss_factor < 1)):
raise Exception(
'`location_loss_factor` should be greater than 0 and less than 1'
)
self.location_loss_factor = location_loss_factor
self._code = code
self.ssd_version = ssd_version
backbone_cut = None
backbone_split = None
if hasattr(self, '_orig_backbone'):
self._backbone_ms = self._backbone
self._backbone = self._orig_backbone
_backbone_meta = cnn_config(self._orig_backbone)
backbone_cut = _backbone_meta['cut']
backbone_split = _backbone_meta['split']
if backbone is None:
self._backbone = models.resnet34
backbone_name = 'res'
elif type(backbone) is str:
self._backbone = getattr(models, backbone)
backbone_name = backbone[:3]
else:
self._backbone = backbone
backbone_name = 'custom'
if not self._check_backbone_support(self._backbone):
raise Exception(
f"Enter only compatible backbones from {', '.join(self.supported_backbones)}"
)
if self._backbone == models.mobilenet_v2:
backbone_cut = -1
backbone_split = _mobilenet_split
if ssd_version == 1:
if grids == None:
grids = [4, 2, 1]
self._create_anchors(grids, zooms, ratios)
feature_sizes = model_sizes(create_body(self._backbone,
cut=backbone_cut),
size=(data.chip_size,
data.chip_size))
num_features = feature_sizes[-1][-1]
num_channels = feature_sizes[-1][1]
ssd_head = SSDHead(grids,
self._anchors_per_cell,
data.c,
num_features=num_features,
drop=drop,
bias=bias,
num_channels=num_channels)
elif ssd_version == 2:
# find bounding boxes height and width
if grids is None:
logger.info("Computing optimal grid size...")
hw = data.height_width
hw = np.array(hw)
# find most suitable centroids for dataset
centroid = kmeans(hw, 1)
avg = avg_iou(hw, centroid)
for num_anchor in range(2, 5):
new_centroid = kmeans(hw, num_anchor)
new_avg = avg_iou(hw, new_centroid)
if (new_avg - avg) < 0.05:
break
avg = new_avg
centroid = new_centroid.copy()
# find grid size
grids = list(
map(
int,
map(
round, data.chip_size /
np.sort(np.max(centroid, axis=1)))))
grids = list(set(grids))
grids.sort(reverse=True)
if grids[-1] == 0:
grids[-1] = 1
grids = list(set(grids))
self._create_anchors(grids, zooms, ratios)
feature_sizes = model_sizes(create_body(self._backbone,
cut=backbone_cut),
size=(data.chip_size,
data.chip_size))
num_features = feature_sizes[-1][-1]
num_channels = feature_sizes[-1][1]
if grids[0] > 8 and abs(num_features - grids[0]
) > 4 and backbone_name == 'res':
num_features = feature_sizes[-2][-1]
num_channels = feature_sizes[-2][1]
backbone_cut = -3
ssd_head = SSDHeadv2(grids,
self._anchors_per_cell,
data.c,
num_features=num_features,
drop=drop,
bias=bias,
num_channels=num_channels)
else:
raise Exception('SSDVersion can only be 1 or 2')
if hasattr(self, '_backbone_ms'):
self._orig_backbone = self._backbone
self._backbone = self._backbone_ms
self.learn = cnn_learner(data=data,
base_arch=self._backbone,
cut=backbone_cut,
split_on=backbone_split,
custom_head=ssd_head)
self._arcgis_init_callback() # make first conv weights learnable
self.learn.model = self.learn.model.to(self._device)
if focal_loss:
self._loss_f = FocalLoss(data.c)
else:
self._loss_f = BCE_Loss(data.c)
self.learn.loss_func = self._ssd_loss
_set_multigpu_callback(self)
if pretrained_path is not None:
self.load(pretrained_path)
def build_res_unet(n_input=1, n_output=2, size=224):
body = create_body(resnet18, pretrained=False, n_in=n_input, cut=-2)
net_G = DynamicUnet(body, n_output, (size, size))
return net_G
def model():
body = create_body(resnet18, pretrained=False)
for param in body.parameters():
param.requires_grad = False
return DynamicUnet(body, 10)
def build_generator(n_inputs, n_outputs, core): body = create_body(core, pretrained=True, n_in=n_inputs, cut=-2) return DynamicUnet(body, n_outputs, (SIZE, SIZE))
def build_res_unet(n_input=1, n_output=2, size=256):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
body = create_body(resnet18, pretrained=True, n_in=n_input, cut=-2)
net_G = DynamicUnet(body, n_output, (size, size)).to(device)
return net_G
