|
- from __future__ import print_function
- import os
- import argparse
- import torch
- import time
- import numpy as np
- import torch.nn as nn
- import torch.nn.functional as F
- import torchvision
- import torch.optim as optim
- from torchvision import datasets, transforms
- from torch.autograd import Variable
- from tiny_imagenet import TinyImageNet
- # from models.imagenet_resnet import resnet50
- from models.RiFT_resnet import *
- from models.RiFT_wideresnet import *
- # from loss.pgd import pgd_loss
- import logging
- from tiny_imagenet import TinyImageNet
-
- os.environ["CUDA_VISIBLE_DEVICES"] = '0'
-
- parser = argparse.ArgumentParser(description='PyTorch CIFAR TRADES Adversarial Training')
- parser.add_argument('--batch-size', type=int, default=128, metavar='N',
- help='input batch size for training (default: 128)')
- parser.add_argument('--test-batch-size', type=int, default=2000, metavar='N',
- help='input batch size for testing (default: 128)')
- parser.add_argument('--data', default='Tiny-ImageNet', help='data')
- parser.add_argument('--arch', default='ResNet18', help='model')
- parser.add_argument('--epochs', type=int, default=100, metavar='N',
- help='number of epochs to train')
- parser.add_argument('--weight-decay', '--wd', default=5e-4,
- type=float, metavar='W')
- parser.add_argument('--lr', type=float, default=0.1, metavar='LR',
- help='learning rate')
- parser.add_argument('--momentum', type=float, default=0.9, metavar='M',
- help='SGD momentum')
- 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=100, metavar='N',
- help='how many batches to wait before logging training status')
- parser.add_argument('--model-dir', default='/tmp/output/baseline/',
- help='directory of model for saving checkpoint')
- parser.add_argument('--save-freq', '-s', default=5, type=int, metavar='N',
- help='save frequency')
- args = parser.parse_args()
-
-
- from c2net.context import prepare,upload_output
-
- #初始化导入数据集和预训练模型到容器内
- c2net_context = prepare()
-
- #获取数据集路径
- tiny_imagenet_200_path = c2net_context.dataset_path#+"/"+"tiny-imagenet-200"
-
- #获取预训练模型路径
- rift_tiny_res_path = c2net_context.pretrain_model_path+"/"
-
- #输出结果必须保存在该目录
- you_should_save_here = c2net_context.output_path
-
-
- # settings
- model_dir = you_should_save_here + '/' + args.arch + '-Standard-' + args.data + '-n'
- if not os.path.exists(model_dir):
- os.makedirs(model_dir)
-
- logger = logging.getLogger(__name__)
- logging.basicConfig(
- format='[%(asctime)s] - %(message)s',
- datefmt='%Y/%m/%d %H:%M:%S',
- level=logging.INFO,
- filename=os.path.join(model_dir, 'nominal_train.log'))
- logger.info(args)
-
-
- use_cuda = not args.no_cuda and torch.cuda.is_available()
- torch.manual_seed(args.seed)
- kwargs = {'num_workers': 4, 'pin_memory': False} if use_cuda else {}
-
- # setup data loader
- if args.data == 'ImageNet':
- transform_train = transforms.Compose([
- transforms.Resize(256),
- transforms.RandomResizedCrop(224),
- transforms.RandomHorizontalFlip(),
- transforms.ToTensor(),
- transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
- ])
- transform_test = transforms.Compose([
- transforms.Resize(256),
- transforms.CenterCrop(224),
- transforms.ToTensor(),
- transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
- ])
- trainset = datasets.ImageFolder('/data/ZNY/data/ImageNet/train', transform_train)
- testset = datasets.ImageFolder("/data/ZNY/data/ImageNet/val", transform_test)
- class_number = 1000
- elif args.data == 'Tiny-ImageNet':
- transform_train = transforms.Compose([
- transforms.RandomCrop(64, padding=8, padding_mode="reflect"),
- transforms.RandomHorizontalFlip(),
- transforms.ToTensor(),
- transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
- ])
- transform_test = transforms.Compose([
- transforms.ToTensor(),
- transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
- ])
- file = zipfile.ZipFile('/tmp/dataset/tiny-imagenet-200.zip', 'r')
- file.extractall('/tmp/dataset/')
- file.close()
- trainset = TinyImageNet('/tmp/dataset/tiny-imagenet-200', train=True, transform=transform_train)
- testset = TinyImageNet('/tmp/dataset/tiny-imagenet-200', train=False, transform=transform_test)
- class_number = 200
- elif args.data == 'CIFAR-10':
- transform_train = transforms.Compose([
- transforms.RandomCrop(32, padding=4),
- transforms.RandomHorizontalFlip(),
- transforms.ToTensor(),
- transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2471, 0.2435, 0.2616)),
- ])
- transform_test = transforms.Compose([
- transforms.ToTensor(),
- transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2471, 0.2435, 0.2616)),
- ])
- trainset = torchvision.datasets.CIFAR10(root=tiny_imagenet_200_path, train=True, download=True, transform=transform_train)
- testset = torchvision.datasets.CIFAR10(root=tiny_imagenet_200_path, train=False, download=True, transform=transform_test)
- class_number = 10
- elif args.data == 'CIFAR-100':
- transform_train = transforms.Compose([
- transforms.RandomCrop(32, padding=4),
- transforms.RandomHorizontalFlip(),
- transforms.ToTensor(),
- transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2471, 0.2435, 0.2616)),
- ])
- transform_test = transforms.Compose([
- transforms.ToTensor(),
- transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2471, 0.2435, 0.2616)),
- ])
- trainset = torchvision.datasets.CIFAR100(root=tiny_imagenet_200_path, train=True, download=True, transform=transform_train)
- testset = torchvision.datasets.CIFAR100(root=tiny_imagenet_200_path, train=False, download=True, transform=transform_test)
- class_number = 100
-
- train_loader = torch.utils.data.DataLoader(trainset, batch_size=args.batch_size, shuffle=True, **kwargs)
- test_loader = torch.utils.data.DataLoader(testset, batch_size=args.test_batch_size, shuffle=False, **kwargs)
-
-
-
- def train(args, model, train_loader, optimizer, epoch):
- model.train()
- for batch_idx, (data, target) in enumerate(train_loader):
- data, target = data.cuda(), target.cuda()
-
- optimizer.zero_grad()
- loss = F.cross_entropy(model(data), target)
- loss.backward()
- optimizer.step()
-
- # print progress
- if batch_idx % args.log_interval == 0:
- print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
- epoch, batch_idx * len(data), len(train_loader.dataset),
- 100. * batch_idx / len(train_loader), loss.item()))
- logger.info('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
- epoch, batch_idx * len(data), len(train_loader.dataset),
- 100. * batch_idx / len(train_loader), loss.item()))
-
-
- def eval_train(model, train_loader):
- model.eval()
- train_loss = 0
- correct = 0
- with torch.no_grad():
- for data, target in train_loader:
- data, target = data.cuda(), target.cuda()
- output = model(data)
- train_loss += F.cross_entropy(output, target, reduction='mean').item()
- pred = output.max(1, keepdim=True)[1]
- correct += pred.eq(target.view_as(pred)).sum().item()
- train_loss /= len(train_loader.dataset)
- print('Training: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)'.format(
- train_loss, correct, len(train_loader.dataset),
- 100. * correct / len(train_loader.dataset)))
- logger.info('Training: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)'.format(
- train_loss, correct, len(train_loader.dataset),
- 100. * correct / len(train_loader.dataset)))
- training_accuracy = correct / len(train_loader.dataset)
- return train_loss, training_accuracy
-
-
- def eval_test(model, test_loader):
- model.eval()
- test_loss = 0
- correct = 0
- with torch.no_grad():
- for data, target in test_loader:
- data, target = data.cuda(), target.cuda()
- output = model(data)
- test_loss += F.cross_entropy(output, target, reduction='mean').item()
- pred = output.max(1, keepdim=True)[1]
- correct += pred.eq(target.view_as(pred)).sum().item()
- test_loss /= len(test_loader.dataset)
- print('Test: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)'.format(
- test_loss, correct, len(test_loader.dataset),
- 100. * correct / len(test_loader.dataset)))
- logger.info('Test: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)'.format(
- test_loss, correct, len(test_loader.dataset),
- 100. * correct / len(test_loader.dataset)))
- test_accuracy = correct / len(test_loader.dataset)
- return test_loss, test_accuracy
-
-
-
-
-
-
- def adjust_learning_rate(optimizer, epoch):
- """decrease the learning rate"""
- lr = args.lr
- if epoch >= 75:
- lr = args.lr * 0.1
- if epoch >= 90:
- lr = args.lr * 0.01
- # if epoch >= 100:
- # lr = args.lr * 0.001
- for param_group in optimizer.param_groups:
- param_group['lr'] = lr
-
-
- def main():
- # init model, ResNet18() can be also used here for training
- if 'CIFAR' in args.data:
- size = 32
- else:
- size = 64
- if args.arch == 'ResNet18':
- model = ResNet18(size, class_number).cuda()
- elif args.arch == 'WideResNet34':
- # model = WideResNet(depth=34, num_classes=class_number, widen_factor=10).cuda()
- model = WideResNet(image_size=size, depth=34, num_classes=class_number, widen_factor=10).cuda()
- elif args.arch == 'ResNet50':
- model = ResNet50(size, class_number).cuda()
-
- # model = torch.nn.DataParallel(model).cuda()
- optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
-
- # test_acc = []
- # time_use = []
-
- for epoch in range(1, args.epochs + 1):
- # adjust learning rate for SGD
- adjust_learning_rate(optimizer, epoch)
-
- start_time = time.time()
-
- # adversarial training
- train(args, model, train_loader, optimizer, epoch)
-
- print('using time:', time.time() - start_time)
- logger.info('using time: {}'.format(time.time() - start_time))
-
-
- print('================================================================')
- logger.info('================================================================')
-
- # evaluation on natural examples
- # train_loss, train_accuracy = eval_train(model, train_loader)
- # test_loss, test_accuracy = eval_test(model, test_loader)
-
- # test_acc.append(test_accuracy)
- # time_use.append(time.time() - start_time)
-
- # file_name = os.path.join(model_dir, 'train_stats.npy')
- # np.save(file_name, np.array(test_acc))
- # file_name = os.path.join(model_dir, 'use_time.npy')
- # np.save(file_name, np.array(time_use))
- torch.save(optimizer.state_dict(),
- os.path.join(model_dir, 'opt-last.tar'))
- torch.save(model.state_dict(),
- os.path.join(model_dir, 'model-last.pth'))
-
- # save checkpoint
- if epoch % args.save_freq == 0:
- torch.save(model.state_dict(),
- os.path.join(model_dir, 'model-{}.pth'.format(epoch)))
- train_loss, train_accuracy = eval_train(model, train_loader)
- test_loss, test_accuracy = eval_test(model, test_loader)
- # 回传结果到openi,只有训练任务才能回传
- upload_output()
-
-
- if __name__ == '__main__':
- main()
|