def test_gl_encdec():
input_x = torch.randn(1, 4, 256, 256)
template_cfg = dict(type='AOTEncoderDecoder')
aot_encdec = build_backbone(template_cfg)
aot_encdec.init_weights()
output = aot_encdec(input_x)
assert output.shape == (1, 3, 256, 256)
cfg_ = template_cfg.copy()
cfg_['encoder'] = dict(type='AOTEncoder')
aot_encdec = build_backbone(cfg_)
output = aot_encdec(input_x)
assert output.shape == (1, 3, 256, 256)
cfg_ = template_cfg.copy()
cfg_['decoder'] = dict(type='AOTDecoder')
aot_encdec = build_backbone(cfg_)
output = aot_encdec(input_x)
assert output.shape == (1, 3, 256, 256)
if torch.cuda.is_available():
aot_encdec = build_backbone(template_cfg)
aot_encdec.init_weights()
aot_encdec = aot_encdec.cuda()
output = aot_encdec(input_x.cuda())
assert output.shape == (1, 3, 256, 256)
def test_gl_encdec():
input_x = torch.randn(1, 4, 256, 256)
template_cfg = dict(type='GLEncoderDecoder')
gl_encdec = build_backbone(template_cfg)
gl_encdec.init_weights()
output = gl_encdec(input_x)
assert output.shape == (1, 3, 256, 256)
cfg_ = template_cfg.copy()
cfg_['decoder'] = dict(type='GLDecoder', out_act='sigmoid')
gl_encdec = build_backbone(cfg_)
output = gl_encdec(input_x)
assert output.shape == (1, 3, 256, 256)
with pytest.raises(ValueError):
cfg_ = template_cfg.copy()
cfg_['decoder'] = dict(type='GLDecoder', out_act='igccc')
gl_encdec = build_backbone(cfg_)
with pytest.raises(TypeError):
gl_encdec.init_weights(pretrained=dict(igccc=4396))
if torch.cuda.is_available():
gl_encdec = build_backbone(template_cfg)
gl_encdec.init_weights()
gl_encdec = gl_encdec.cuda()
output = gl_encdec(input_x.cuda())
assert output.shape == (1, 3, 256, 256)
def test_dic_net():
model_cfg = dict(
type='DICNet',
in_channels=3,
out_channels=3,
mid_channels=48,
num_blocks=6,
hg_mid_channels=256,
hg_num_keypoints=68,
num_steps=4,
upscale_factor=8,
detach_attention=False)
# build model
model = build_backbone(model_cfg)
# test attributes
assert model.__class__.__name__ == 'DICNet'
# prepare data
inputs = torch.rand(1, 3, 16, 16)
targets = torch.rand(1, 3, 128, 128)
# prepare loss
loss_function = nn.L1Loss()
# prepare optimizer
optimizer = torch.optim.Adam(model.parameters())
# test on cpu
output, _ = model(inputs)
optimizer.zero_grad()
loss = loss_function(output[-1], targets)
loss.backward()
optimizer.step()
assert len(output) == 4
assert torch.is_tensor(output[-1])
assert output[-1].shape == targets.shape
# test on gpu
if torch.cuda.is_available():
model = model.cuda()
optimizer = torch.optim.Adam(model.parameters())
inputs = inputs.cuda()
targets = targets.cuda()
output, _ = model(inputs)
optimizer.zero_grad()
loss = loss_function(output[-1], targets)
loss.backward()
optimizer.step()
assert len(output) == 4
assert torch.is_tensor(output[-1])
assert output[-1].shape == targets.shape
with pytest.raises(OSError):
model.init_weights('')
with pytest.raises(TypeError):
model.init_weights(1)
def test_rdn():
scale = 4
model_cfg = dict(
type='RDN',
in_channels=3,
out_channels=3,
mid_channels=64,
num_blocks=16,
upscale_factor=scale)
# build model
model = build_backbone(model_cfg)
# test attributes
assert model.__class__.__name__ == 'RDN'
# prepare data
inputs = torch.rand(1, 3, 32, 16)
targets = torch.rand(1, 3, 128, 64)
# prepare loss
loss_function = nn.L1Loss()
# prepare optimizer
optimizer = torch.optim.Adam(model.parameters())
# test on cpu
output = model(inputs)
optimizer.zero_grad()
loss = loss_function(output, targets)
loss.backward()
optimizer.step()
assert torch.is_tensor(output)
assert output.shape == targets.shape
# test on gpu
if torch.cuda.is_available():
model = model.cuda()
optimizer = torch.optim.Adam(model.parameters())
inputs = inputs.cuda()
targets = targets.cuda()
output = model(inputs)
optimizer.zero_grad()
loss = loss_function(output, targets)
loss.backward()
optimizer.step()
assert torch.is_tensor(output)
assert output.shape == targets.shape
def test_liif_rdn():
model_cfg = dict(type='LIIFRDN',
encoder=dict(type='RDN',
in_channels=3,
out_channels=3,
mid_channels=64,
num_blocks=16,
upscale_factor=4,
num_layers=8,
channel_growth=64),
imnet=dict(type='MLPRefiner',
in_dim=64,
out_dim=3,
hidden_list=[256, 256, 256, 256]),
local_ensemble=True,
feat_unfold=True,
cell_decode=True,
eval_bsize=30000)
# build model
model = build_backbone(model_cfg)
# test attributes
assert model.__class__.__name__ == 'LIIFRDN'
# prepare data
inputs = torch.rand(1, 3, 22, 11)
targets = torch.rand(1, 128 * 64, 3)
coord = torch.rand(1, 128 * 64, 2)
cell = torch.rand(1, 128 * 64, 2)
# test on cpu
output = model(inputs, coord, cell)
output = model(inputs, coord, cell, True)
assert torch.is_tensor(output)
assert output.shape == targets.shape
# test on gpu
if torch.cuda.is_available():
model = model.cuda()
inputs = inputs.cuda()
targets = targets.cuda()
coord = coord.cuda()
cell = cell.cuda()
output = model(inputs, coord, cell)
output = model(inputs, coord, cell, True)
assert torch.is_tensor(output)
assert output.shape == targets.shape
def test_ttsr_net():
inputs = torch.rand(2, 3, 24, 24)
soft_attention = torch.rand(2, 1, 24, 24)
t_level3 = torch.rand(2, 64, 24, 24)
t_level2 = torch.rand(2, 32, 48, 48)
t_level1 = torch.rand(2, 16, 96, 96)
ttsr_cfg = dict(type='TTSRNet',
in_channels=3,
out_channels=3,
mid_channels=16,
texture_channels=16)
ttsr = build_backbone(ttsr_cfg)
outputs = ttsr(inputs, soft_attention, (t_level3, t_level2, t_level1))
assert outputs.shape == (2, 3, 96, 96)
def test_unet_generator():
# color to color
cfg = dict(type='UnetGenerator',
in_channels=3,
out_channels=3,
num_down=8,
base_channels=64,
norm_cfg=dict(type='BN'),
use_dropout=True,
init_cfg=dict(type='normal', gain=0.02))
net = build_backbone(cfg)
net.init_weights(pretrained=None)
# cpu
input_shape = (1, 3, 256, 256)
img = _demo_inputs(input_shape)
output = net(img)
assert output.shape == (1, 3, 256, 256)
# gpu
if torch.cuda.is_available():
net = net.cuda()
output = net(img.cuda())
assert output.shape == (1, 3, 256, 256)
# gray to color
cfg = dict(type='UnetGenerator',
in_channels=1,
out_channels=3,
num_down=8,
base_channels=64,
norm_cfg=dict(type='BN'),
use_dropout=True,
init_cfg=dict(type='normal', gain=0.02))
net = build_backbone(cfg)
net.init_weights(pretrained=None)
# cpu
input_shape = (1, 1, 256, 256)
img = _demo_inputs(input_shape)
output = net(img)
assert output.shape == (1, 3, 256, 256)
# gpu
if torch.cuda.is_available():
net = net.cuda()
output = net(img.cuda())
assert output.shape == (1, 3, 256, 256)
# color to gray
cfg = dict(type='UnetGenerator',
in_channels=3,
out_channels=1,
num_down=8,
base_channels=64,
norm_cfg=dict(type='BN'),
use_dropout=True,
init_cfg=dict(type='normal', gain=0.02))
net = build_backbone(cfg)
net.init_weights(pretrained=None)
# cpu
input_shape = (1, 3, 256, 256)
img = _demo_inputs(input_shape)
output = net(img)
assert output.shape == (1, 1, 256, 256)
# gpu
if torch.cuda.is_available():
net = net.cuda()
output = net(img.cuda())
assert output.shape == (1, 1, 256, 256)
# pretrained should be str or None
with pytest.raises(TypeError):
net.init_weights(pretrained=[1])
# test norm_cfg assertions
bad_cfg = copy.deepcopy(cfg)
bad_cfg['norm_cfg'] = None
with pytest.raises(AssertionError):
_ = build_backbone(bad_cfg)
bad_cfg['norm_cfg'] = dict(tp='BN')
with pytest.raises(AssertionError):
_ = build_backbone(bad_cfg)
def test_resnet_generator():
# color to color
cfg = dict(type='ResnetGenerator',
in_channels=3,
out_channels=3,
base_channels=64,
norm_cfg=dict(type='IN'),
use_dropout=False,
num_blocks=9,
padding_mode='reflect',
init_cfg=dict(type='normal', gain=0.02))
net = build_backbone(cfg)
net.init_weights(pretrained=None)
# cpu
input_shape = (1, 3, 256, 256)
img = _demo_inputs(input_shape)
output = net(img)
assert output.shape == (1, 3, 256, 256)
# gpu
if torch.cuda.is_available():
net = net.cuda()
output = net(img.cuda())
assert output.shape == (1, 3, 256, 256)
# gray to color
cfg = dict(type='ResnetGenerator',
in_channels=1,
out_channels=3,
base_channels=64,
norm_cfg=dict(type='IN'),
use_dropout=False,
num_blocks=9,
padding_mode='reflect',
init_cfg=dict(type='normal', gain=0.02))
net = build_backbone(cfg)
net.init_weights(pretrained=None)
# cpu
input_shape = (1, 1, 256, 256)
img = _demo_inputs(input_shape)
output = net(img)
assert output.shape == (1, 3, 256, 256)
# gpu
if torch.cuda.is_available():
net = net.cuda()
output = net(img.cuda())
assert output.shape == (1, 3, 256, 256)
# color to gray
cfg = dict(type='ResnetGenerator',
in_channels=3,
out_channels=1,
base_channels=64,
norm_cfg=dict(type='IN'),
use_dropout=False,
num_blocks=9,
padding_mode='reflect',
init_cfg=dict(type='normal', gain=0.02))
net = build_backbone(cfg)
net.init_weights(pretrained=None)
# cpu
input_shape = (1, 3, 256, 256)
img = _demo_inputs(input_shape)
output = net(img)
assert output.shape == (1, 1, 256, 256)
# gpu
if torch.cuda.is_available():
net = net.cuda()
output = net(img.cuda())
assert output.shape == (1, 1, 256, 256)
# test num_blocks non-negative
bad_cfg = copy.deepcopy(cfg)
bad_cfg['num_blocks'] = -1
with pytest.raises(AssertionError):
net = build_backbone(bad_cfg)
# pretrained should be str or None
with pytest.raises(TypeError):
net.init_weights(pretrained=[1])
# test norm_cfg assertions
bad_cfg = copy.deepcopy(cfg)
bad_cfg['norm_cfg'] = None
with pytest.raises(AssertionError):
_ = build_backbone(bad_cfg)
bad_cfg['norm_cfg'] = dict(tp='IN')
with pytest.raises(AssertionError):
_ = build_backbone(bad_cfg)
def test_cain_net():
model_cfg = dict(type='CAINNet')
# build model
model = build_backbone(model_cfg)
# test attributes
assert model.__class__.__name__ == 'CAINNet'
# prepare data
inputs0 = torch.rand(1, 2, 3, 5, 5)
target0 = torch.rand(1, 3, 5, 5)
inputs = torch.rand(1, 2, 3, 256, 248)
target = torch.rand(1, 3, 256, 248)
# test on cpu
output = model(inputs)
output = model(inputs, padding_flag=True)
model(inputs0, padding_flag=True)
assert torch.is_tensor(output)
assert output.shape == target.shape
with pytest.raises(AssertionError):
output = model(inputs[:, :1])
with pytest.raises(OSError):
model.init_weights('')
with pytest.raises(TypeError):
model.init_weights(1)
model_cfg = dict(type='CAINNet', norm='in')
model = build_backbone(model_cfg)
model(inputs)
model_cfg = dict(type='CAINNet', norm='bn')
model = build_backbone(model_cfg)
model(inputs)
with pytest.raises(ValueError):
model_cfg = dict(type='CAINNet', norm='lys')
build_backbone(model_cfg)
# test on gpu
if torch.cuda.is_available():
model = model.cuda()
inputs = inputs.cuda()
target = target.cuda()
output = model(inputs)
output = model(inputs, True)
assert torch.is_tensor(output)
assert output.shape == target.shape
inputs0 = inputs0.cuda()
target0 = target0.cuda()
model(inputs0, padding_flag=True)
model_cfg = dict(type='CAINNet', norm='in')
model = build_backbone(model_cfg).cuda()
model(inputs)
model_cfg = dict(type='CAINNet', norm='bn')
model = build_backbone(model_cfg).cuda()
model(inputs)
with pytest.raises(ValueError):
model_cfg = dict(type='CAINNet', norm='lys')
build_backbone(model_cfg).cuda()