def test_lstm():
results = summary(LSTMNet(), (100, ), dtypes=[torch.long])
assert len(results.summary_list) == 3, "Should find 3 layers"
def test_exceptions():
with pytest.raises(RuntimeError):
summary(EdgeCaseModel(throw_error=True), (1, 28, 28))
with pytest.raises(TypeError):
summary(EdgeCaseModel(return_str=True), (1, 28, 28))
def test_input_tensor():
input_data = torch.randn(5, 1, 28, 28)
metrics = summary(SingleInputNet(), input_data)
assert metrics.input_size == [torch.Size([5, 1, 28, 28])]
def test_functional_layers():
summary(FunctionalNet(), (1, 28, 28))
def test_return_dict():
input_size = [torch.Size([1, 28, 28]), [12]]
metrics = summary(ReturnDict(), input_size, col_width=65)
assert metrics.input_size == [(1, 28, 28), [12]]
def test_input_size_possibilities():
test = CustomModule(2, 3)
summary(test, [(2, )])
summary(test, ((2, ), ))
summary(test, (2, ))
summary(test, [2])
with pytest.raises(AssertionError):
summary(test, [(3, 0)])
with pytest.raises(TypeError):
summary(test, {0: 1})
with pytest.raises(TypeError):
summary(test, "hello")
def test_siamese_net():
metrics = summary(SiameseNets(), [(1, 88, 88), (1, 88, 88)])
assert round(metrics.to_megabytes(metrics.total_input), 2) == 0.06
def test_model_with_args():
summary(RecursiveNet(), (64, 28, 28), "args1", args2="args2")
without_t_stride=without_t_stride,
temporal_module=temporal_module, dropout=dropout, fpn_dim=fpn_dim,
pooling_method=pooling_method, input_channels=input_channels)
if add_non_local:
model = make_non_local_v2(model, groups)
if imagenet_pretrained:
state_dict = model_zoo.load_url(model_urls['resnet{}'.format(depth)], map_location='cpu')
state_dict.pop('fc.weight', None)
state_dict.pop('fc.bias', None)
if input_channels != 3: # convert the RGB model to others, like flow
state_dict = convert_rgb_model_to_others(state_dict, input_channels, 7)
model.load_state_dict(state_dict, strict=False)
return model
if __name__ == '__main__':
from torchsummary import torchsummary
model = resnet(18, 187, without_t_stride=True, groups=8, dropout=0.5,
temporal_module_name=None, dw_conv=True, blending_frames=3,
blending_method='sum', fpn_dim=-1, pooling_method='max', input_channels=3)
dummy_data = (24, 224, 224)
model.eval()
model_summary = torchsummary.summary(model, input_size=dummy_data)
# print(model_summary)
print(model)
print(model.network_name)
def train(args):
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
init_net = nn.DataParallel(models.InitNet(args)).to(device)
deep_net = nn.DataParallel(models.UNet(args)).to(device)
print("Data loading.")
dataset = utils.loader(args)
print("Data loaded.")
criterion = nn.L1Loss().to(device)
optimizer_init = optim.Adam(init_net.parameters())
optimizer_deep = optim.Adam(deep_net.parameters())
scheduler_init = optim.lr_scheduler.MultiStepLR(optimizer_init, milestones=[50, 80], gamma=0.1)
scheduler_deep = optim.lr_scheduler.MultiStepLR(optimizer_deep, milestones=[50, 80], gamma=0.1)
print("Train start.")
time_start = time.time()
if os.path.exists(args.init_state_dict) and os.path.exists(args.deep_state_dict):
if torch.cuda.is_available():
checkpoint_init = torch.load(args.init_state_dict)
checkpoint_deep = torch.load(args.deep_state_dict)
else:
checkpoint_init = torch.load(args.init_state_dict, map_location="cpu")
checkpoint_deep = torch.load(args.deep_state_dict, map_location="cpu")
init_net.load_state_dict(checkpoint_init["model"])
optimizer_init.load_state_dict(checkpoint_init["optimizer"])
deep_net.load_state_dict(checkpoint_deep["model"])
optimizer_deep.load_state_dict(checkpoint_deep["optimizer"])
start_epoch = checkpoint_deep["epoch"]
print("Success loading epoch {}".format(start_epoch))
else:
start_epoch = 0
print("No saved model, start epoch = 0.")
for epoch in range(start_epoch, args.epochs):
for idx, item in enumerate(dataset):
x, _ = item
x = x.to(device)
optimizer_init.zero_grad()
optimizer_deep.zero_grad()
init_x = init_net(x)
init_x = utils.reshape(init_x, args)
deep_x = deep_net(init_x)
loss_init = criterion(x, init_x)
loss_deep = criterion(x, init_x + deep_x)
loss_init.backward(retain_graph=True)
loss_deep.backward()
optimizer_init.step()
optimizer_deep.step()
use_time = time.time() - time_start
if (idx + 1) % 20 == 0:
print("=> epoch: {}, batch: {}, Loss1: {:.4f}, Loss2: {:.4f}, lr1: {}, lr2: {}, used time: {:.4f}"
.format(epoch + 1, idx + 1, loss_init.item(), loss_deep.item(),
optimizer_init.param_groups[0]['lr'], optimizer_deep.param_groups[0]['lr'], use_time))
scheduler_init.step()
scheduler_deep.step()
state_init = {"model": init_net.state_dict(), "optimizer": optimizer_init.state_dict()}
state_deep = {"model": deep_net.state_dict(), "optimizer": optimizer_deep.state_dict(), "epoch": epoch + 1}
torch.save(state_init, args.init_state_dict)
torch.save(state_deep, args.deep_state_dict)
print("Check point of epoch {} saved.".format(epoch + 1))
print("Train end.")
torchsummary.summary(init_net, (1, 32, 32))
torchsummary.summary(deep_net, (1, 32, 32))
with open("./trained_models/init_net.txt", "w") as f1:
f1.write(torchsummary.summary(init_net, (1, 32, 32)))
with open("./trained_models/deep_net.txt", "w") as f2:
f2.write(torchsummary.summary(deep_net, (1, 32, 32)))
def test_exception_output(capsys):
summary(EdgeCaseModel(throw_error=False), (1, 28, 28))
with pytest.raises(RuntimeError):
summary(EdgeCaseModel(throw_error=True), (1, 28, 28))
verify_output(capsys, "unit_test/test_output/exception.out")
def test_resnet_out(capsys):
model = torchvision.models.resnet152()
summary(model, (3, 224, 224), depth=3)
verify_output(capsys, "unit_test/test_output/resnet152.out")
