def get_model(dataset_type):
if dataset_type == 'mnist':
model = MNISTNet().to(device)
elif dataset_type == 'cifar10':
# model = CIFAR10Net().to(device)
model = CNN9Layer(num_classes=10, input_shape=3).to(device)
# model = ResNet18(num_classes=10).to(device)
elif dataset_type == 'cifar100':
# model = CNN9Layer(num_classes=100, input_shape=3).to(device)
model = ResNet18(num_classes=100).to(device)
return model
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--config', type=str, default='configs/config.json')
parser.add_argument('--no-cuda', action='store_true')
parser.add_argument('--parallel', action='store_true')
args = parser.parse_args()
args.cuda = torch.cuda.is_available() and not args.no_cuda
print(args)
device = torch.device('cuda' if args.cuda else 'cpu')
config = load_json(args.config)
model = MNISTNet()
if args.parallel:
model = nn.DataParallel(model)
model.to(device)
optimizer = optim.Adam(model.parameters(), **config['adam'])
scheduler = optim.lr_scheduler.StepLR(optimizer, **config['steplr'])
train_loader, valid_loader = mnist_loader(**config['dataset'])
trainer = Trainer(model, optimizer, train_loader, valid_loader, device)
output_dir = os.path.join(config['output_dir'],
datetime.now().strftime('%Y%m%d_%H%M%S'))
os.makedirs(output_dir, exist_ok=True)
# save config to output dir
save_json(config, os.path.join(output_dir, 'config.json'))
for epoch in range(config['epochs']):
scheduler.step()
train_loss, train_acc = trainer.train()
valid_loss, valid_acc = trainer.validate()
print(
'epoch: {}/{},'.format(epoch + 1, config['epochs']),
'train loss: {:.4f}, train acc: {:.2f}%,'.format(
train_loss, train_acc * 100),
'valid loss: {:.4f}, valid acc: {:.2f}%'.format(
valid_loss, valid_acc * 100))
torch.save(
model.state_dict(),
os.path.join(output_dir, 'model_{:04d}.pt'.format(epoch + 1)))
log_path = common.create_log_dir(args.exp_id, task, args.loss)
print(f"Logging to {log_path}")
# Dumping all script arguments
common.dump_params(join(log_path, 'config.cfg'), args)
# Set custom seed before doing anything
common.set_custom_seed(args.seed)
# Load dataset and create model
print(f"[Task: {task.upper()}]")
print(f"[Loss: {args.loss.upper()}]")
print('[Loading Dataset...]')
nfeat, nclass = 2, 10
config = common.get_config(args.loss, nfeat, nclass, task, args.margin)
model = MNISTNet(nfeat, loss_module=config.loss_module)
dataset = MNIST(args.path, args.batch_size)
dev = dataset.dev_partition()
train = dataset.training_partition()
print('[Dataset Loaded]')
# Train and evaluation plugins
test_callbacks = []
train_callbacks = []
# Logging configuration
if args.log_interval in range(1, 101):
print(
f"[Logging: {common.enabled_str(True)} (every {args.log_interval}%)]")
if __name__ == "__main__":
set_seed(args.seed)
# create output folder
now = time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime())
output_path = os.path.join(args.log_dir, args.project_name + ' ' + now)
os.makedirs(output_path)
# get data loader
train_loader, val_loader, test_loader, clean_sample_idx, noisy_sample_idx, dataset_len = create_dataloader(
args.dataset, args.dataset_path, args.noise_type, args.noise_rate)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if args.dataset == 'mnist':
model = MNISTNet().to(device)
elif args.dataset == 'cifar10':
# model = CIFAR10Net().to(device)
model = CNN9Layer(num_classes=10, input_shape=3).to(device)
# model = ResNet18(num_classes=10).to(device)
elif args.dataset == 'cifar100':
# model = CNN9Layer(num_classes=100, input_shape=3).to(device)
model = ResNet18(num_classes=100).to(device)
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=0.9,
weight_decay=args.l2_reg)
train_loss_lst, val_loss_lst = [], []
train_acc_lst, val_acc_lst = [], []
def main():
parser = argparse.ArgumentParser(
description='PyTorch L-Softmax MNIST Example')
parser.add_argument(
'--margin',
type=int,
default=4,
metavar='M',
help='the margin for the l-softmax formula (m=1, 2, 3, 4)')
parser.add_argument('--batch-size',
type=int,
default=256,
metavar='N',
help='input batch size for training (default: 256)')
parser.add_argument('--test-batch-size',
type=int,
default=1000,
metavar='N',
help='input batch size for testing (default: 1000)')
parser.add_argument('--epochs',
type=int,
default=100,
metavar='N',
help='number of epochs to train (default: 100)')
parser.add_argument('--lr',
type=float,
default=0.1,
metavar='LR',
help='initial learning rate (default: 0.1)')
parser.add_argument('--momentum',
type=float,
default=0.9,
metavar='M',
help='SGD momentum (default: 0.9)')
parser.add_argument('--weight-decay',
type=float,
default=5e-4,
metavar='W',
help='SGD weight decay (default: 0.0005)')
parser.add_argument('--no-cuda',
action='store_true',
default=False,
help='disables CUDA training')
parser.add_argument('--seed',
type=int,
default=1,
metavar='S',
help='random seed (default: 1)')
parser.add_argument(
'--log-interval',
type=int,
default=10,
metavar='N',
help='how many batches to wait before logging training status')
parser.add_argument(
'--vis',
default=False,
metavar='V',
help='enables visualizing 2d features (default: False).')
args = parser.parse_args()
print(args)
use_cuda = not args.no_cuda and torch.cuda.is_available()
torch.manual_seed(args.seed)
device = torch.device("cuda" if use_cuda else "cpu")
kwargs = {'num_workers': 4, 'pin_memory': True} if use_cuda else {}
train_loader = torch.utils.data.DataLoader(datasets.MNIST(
'../data',
train=True,
download=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307, ), (0.3081, ))
])),
batch_size=args.batch_size,
shuffle=True,
**kwargs)
test_loader = torch.utils.data.DataLoader(datasets.MNIST(
'../data',
train=False,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307, ), (0.3081, ))
])),
batch_size=args.test_batch_size,
shuffle=True,
**kwargs)
criterion = nn.CrossEntropyLoss().to(device)
if args.vis:
model = MNISTFIG2Net(margin=args.margin, device=device).to(device)
else:
model = MNISTNet(margin=args.margin, device=device).to(device)
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum)
for epoch in range(1, args.epochs + 1):
adjust_learning_rate(args, optimizer, epoch)
train(args, model, criterion, device, train_loader, optimizer, epoch)
test(args, model, criterion, device, test_loader)
if args.vis:
plot_2d_features(args, model, device, test_loader)
# create output folder
now = time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime())
output_path = os.path.join(args.log_dir, args.project_name + ' ' + now)
os.makedirs(output_path)
# get data loader
train_loader, val_loader, test_loader, clean_sample_idx, noisy_sample_idx, dataset_len = create_dataloader(
args.dataset,
args.dataset_path,
args.noise_type,
args.noise_rate)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if args.dataset == 'mnist':
model = MNISTNet().to(device)
elif args.dataset == 'cifar10':
model = CNN9Layer(num_classes=args.num_classes, input_shape=1).to(device)
optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=0.9, weight_decay=args.l2_reg)
train_loss_lst, val_loss_lst = [], []
train_acc_lst, val_acc_lst = [], []
clean_mean_loss_lst, noisy_mean_loss_lst = [], []
clean_min_loss_lst, noisy_min_loss_lst = [], []
clean_max_loss_lst, noisy_max_loss_lst = [], []
sample_loss = np.zeros(dataset_len)
# main loop(train,val,test)