5df865b7f19376e82ce9b93688cc8a877770a51c,imagenet.py,,main,#Any#Any#,30

Before Change

def main(batch_size, root):
    train_loader, test_loader = get_dataloader(batch_size, root)
    _se_resnet = se_resnet50(num_classes=1000)
    se_resnet = nn.DataParallel(_se_resnet, device_ids=[0, 1])
    optimizer = optim.SGD(params=se_resnet.parameters(), lr=0.6, momentum=0.9, weight_decay=1e-4)
    scheduler = StepLR(optimizer, 30, gamma=0.1)
    trainer = Trainer(se_resnet, optimizer, F.cross_entropy, save_dir=".")

After Change

def main(batch_size, root):
    train_loader, test_loader = get_dataloader(batch_size, root)
    se_resnet = nn.DataParallel(se_resnet50(num_classes=1000),
                                device_ids=list(range(torch.cuda.device_count())))
    optimizer = optim.SGD(params=se_resnet.parameters(), lr=0.6, momentum=0.9, weight_decay=1e-4)
    scheduler = optim.lr_scheduler.StepLR(optimizer, 30, gamma=0.1)
    trainer = Trainer(se_resnet, optimizer, F.cross_entropy, save_dir=".")
    trainer.loop(100, train_loader, test_loader, scheduler)

Instances