ca078e83147f95ec7df149ef5eea9b6a8424d3bf,memcnn/models/tests/test_memory_saving.py,,test_memory_saving,#,199

Before Change

    network.to(device)
    network.train()
    network.zero_grad()
    optim = torch.optim.RMSprop(network.parameters())
    optim.zero_grad()
    mem_start = 0 if not device == "cuda" else \
        torch.cuda.memory_allocated() / float(1024 ** 2)

    y = network(xx)
    gc.collect()
    mem_after_forward = mem_reporter.collect_stats() / float(1024 ** 2) if not device == "cuda" else \
        torch.cuda.memory_allocated() / float(1024 ** 2)
    loss = torch.nn.MSELoss()(y, ytarget)
    optim.zero_grad()
    loss.backward()
    optim.step()
    gc.collect()
    // mem_after_backward = mem_reporter.collect_stats() / float(1024 ** 2) if not device == "cuda" else \
    //     torch.cuda.memory_allocated() / float(1024 ** 2)
    gc.enable()

After Change

        network.to(device)
        network.train()
        network.zero_grad()
        optim = torch.optim.RMSprop(network.parameters())
        optim.zero_grad()
        mem_start = 0 if not device == "cuda" else \
            torch.cuda.memory_allocated() / float(1024 ** 2)

        y = network(xx)
        gc.collect()
        mem_after_forward = mem_reporter.collect_stats() / float(1024 ** 2) if not device == "cuda" else \
            torch.cuda.memory_allocated() / float(1024 ** 2)
        loss = torch.nn.MSELoss()(y, ytarget)
        optim.zero_grad()
        loss.backward()
        optim.step()
        gc.collect()
        // mem_after_backward = mem_reporter.collect_stats() / float(1024 ** 2) if not device == "cuda" else \
        //     torch.cuda.memory_allocated() / float(1024 ** 2)
        gc.enable()

Instances