def main(args): net = smallnet(in_channels=1, num_targets=10) augerino = models.UniformAug() model = models.AugAveragedModel(net, augerino, ncopies=args.ncopies) start_widths = torch.ones(6) * -5. start_widths[2] = 1. model.aug.set_width(start_widths) softplus = torch.nn.Softplus() dataset = datasets.RotMNIST("~/datasets/", train=True) trainloader = DataLoader(dataset, batch_size=args.batch_size) optimizer = torch.optim.Adam([{ 'name': 'model', 'params': model.model.parameters(), "weight_decay": args.wd }, { 'name': 'aug', 'params': model.aug.parameters(), "weight_decay": 0. }], lr=args.lr) use_cuda = torch.cuda.is_available() if use_cuda: model = model.cuda() print("Using Cuda") ## save init model ## fname = "/model" + str(args.aug_reg) + "_init.pt" torch.save(model.state_dict(), args.dir + fname) criterion = losses.safe_unif_aug_loss logger = [] for epoch in range(args.epochs): # loop over the dataset multiple times epoch_loss = 0 batches = 0 for i, data in enumerate(trainloader, 0): # get the inputs; data is a list of [inputs, labels] inputs, labels = data if use_cuda: inputs, labels = inputs.cuda(), labels.cuda() # zero the parameter gradients optimizer.zero_grad() # forward + backward + optimize # print(inputs.shape) outputs = model(inputs) loss = criterion(outputs, labels, model, reg=args.aug_reg) loss.backward() optimizer.step() epoch_loss += loss.detach().item() batches += 1 print(epoch, loss.item(), softplus(model.aug.width).detach().data) log = model.aug.width.tolist() log += model.aug.width.grad.data.tolist() log += [loss.item()] logger.append(log) fname = "/model" + str(args.aug_reg) + ".pt" torch.save(model.state_dict(), args.dir + fname) df = pd.DataFrame(logger) df.to_pickle(args.dir + "/auglog_" + str(args.aug_reg) + ".pkl")
def main(args): net = models.layer13s(in_channels=3, num_targets=10) augerino = models.UniformAug() model = models.AugAveragedModel(net, augerino, ncopies=args.ncopies) transform = transforms.Compose([ # you can add other transformations in this list transforms.ToTensor() ]) fname = "aug_reg" + str(args.aug_reg) + ".pt" dataset = torchvision.datasets.CIFAR10("/datasets/cifar10", train=True, download=False, transform=transform) trainloader = DataLoader(dataset, batch_size=args.batch_size) optimizer = torch.optim.Adam([{ 'name': 'model', 'params': model.model.parameters(), "weight_decay": args.wd }, { 'name': 'aug', 'params': model.aug.parameters(), "weight_decay": 0. }], lr=args.lr) use_cuda = torch.cuda.is_available() if use_cuda: model = model.cuda() print("Using Cuda") criterion = losses.safe_unif_aug_loss for epoch in range(args.epochs): # loop over the dataset multiple times epoch_loss = 0 batches = 0 for i, data in enumerate(trainloader, 0): # get the inputs; data is a list of [inputs, labels] inputs, labels = data if use_cuda: inputs, labels = inputs.cuda(), labels.cuda() # zero the parameter gradients optimizer.zero_grad() # forward + backward + optimize # print(inputs.shape) outputs = model(inputs) loss = criterion(outputs, labels, model, reg=args.aug_reg) loss.backward() optimizer.step() epoch_loss += loss.item() batches += 1 print("Epoch = ", epoch) print("Epoch loss = ", epoch_loss / batches) print("\n") torch.save(model.state_dict(), args.dir + fname)
def main(args): net = models.layer13s(in_channels=3,num_targets=10) augerino = models.UniformAug() model = models.AugAveragedModel(net, augerino,ncopies=args.ncopies) if args.transform == 'none': transform = transforms.Compose([ # you can add other transformations in this list transforms.ToTensor() ]) root_name = "aug_no_trans_" else: transform = transforms.Compose([ transforms.RandomHorizontalFlip(), transforms.Resize(32), transforms.RandomCrop(32, padding=4), transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) ]) root_name = "aug_fixed_trans_" dataset = torchvision.datasets.CIFAR10(args.data_dir, train=True, download=False, transform=transform) trainloader = DataLoader(dataset, batch_size=args.batch_size) optimizer = torch.optim.Adam([{'name': 'model', 'params': model.model.parameters(), "weight_decay": args.wd}, {'name': 'aug', 'params': model.aug.parameters(), "weight_decay": 0.}], lr=args.lr) use_cuda = torch.cuda.is_available() if use_cuda: model = model.cuda() print("Using Cuda") ## save init model ## fname = root_name + "init.pt" torch.save(model.state_dict(), args.dir + fname) criterion = losses.safe_unif_aug_loss for epoch in range(args.epochs): # loop over the dataset multiple times epoch_loss = 0 batches = 0 for i, data in enumerate(trainloader, 0): # get the inputs; data is a list of [inputs, labels] inputs, labels = data if use_cuda: inputs, labels = inputs.cuda(), labels.cuda() # zero the parameter gradients optimizer.zero_grad() # forward + backward + optimize # print(inputs.shape) outputs = model(inputs) loss = criterion(outputs, labels, model, reg=args.aug_reg) loss.backward() optimizer.step() epoch_loss += loss.item() batches += 1 print("Epoch = ", epoch) print("Epoch loss = ", epoch_loss/batches) print("\n") fname = root_name + "_trained.pt" torch.save(model.state_dict(), args.dir + fname)