def build_trainer(self, config, args):
if not os.path.exists(args.model_save_path):
_logger.warn("{} not exists, create it".format(
args.model_save_path))
os.makedirs(args.model_save_path)
_set_file(args.model_save_path + 'log.log')
import torchvision.transforms as transforms
CIFAR_MEAN = [0.49139968, 0.48215827, 0.44653124]
CIFAR_STD = [0.24703233, 0.24348505, 0.26158768]
train_transform = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(CIFAR_MEAN, CIFAR_STD),
])
train_data = dset.CIFAR10(root='/data/limingyao/.torch/datasets',
train=True,
download=False,
transform=train_transform)
num_train = len(train_data)
indices = list(range(num_train))
split = int(np.floor(config.train_portion * num_train))
train_queue = torch.utils.data.DataLoader(train_data,
batch_size=args.batch_size,
sampler=SubsetRandomSampler(
indices[:split]),
pin_memory=True,
num_workers=16)
val_queue = torch.utils.data.DataLoader(train_data,
batch_size=args.batch_size,
sampler=SubsetRandomSampler(
indices[split:]),
pin_memory=True,
num_workers=8)
self.train_queue = train_queue
self.val_queue = val_queue
model = EDNet(
num_classes=config.num_cls_used if config.num_cls_used > 0 else 10,
width_mult=config.width_mult,
kernel_size=config.kernel_size)
self.trainer = Trainer(network=model,
w_lr=config.w_lr,
w_mom=config.w_mom,
w_wd=config.w_wd,
init_temperature=config.init_temperature,
temperature_decay=config.temperature_decay,
logger=_logger,
lr_scheduler=lr_scheduler_params,
target_lat=config.target_lat,
gpus=args.gpus,
model_save_path=args.model_save_path)
help='log frequence, default is 100')
parser.add_argument('--patch-idx', type=int, default=0,
help='patch index, default is 0')
parser.add_argument('--patch-size', type=int, default=1,
help='patch size, default is 1')
parser.add_argument('--gpus', type=str, default='0',
help='gpus, default is 0')
parser.add_argument('--load-model-path', type=str, default=None,
help='re_train, default is None')
args = parser.parse_args()
args.model_save_path = '/data/limingyao/model/nas/snas/%s/' % \
(time.strftime('%Y-%m-%d', time.localtime(time.time())))
if not os.path.exists(args.model_save_path):
_logger.warn("{} not exists, create it".format(args.model_save_path))
os.makedirs(args.model_save_path)
_set_file(args.model_save_path + 'log.log')
config = Config()
train_transform, valid_transform = utils._data_transforms_cifar10(config)
train_data = dset.CIFAR10(root="/data/limingyao/.torch/datasets",train=True,
download=True, transform=train_transform)
train_queue = torch.utils.data.DataLoader(
train_data, batch_size=args.batch_size,
shuffle=True, pin_memory=True, num_workers=16)
model = SNAS(C=config.init_channels,
num_classes=10,
layers=config.stacked_cell)
