default=None,
help='load pretrained checkpoint for assigning labels')
parser.add_argument('--pretrain_epo',
type=int,
default=None,
help='load pretrained epoch for assigning labels')
args = parser.parse_args()
## Pipelines Setup
model_dir = os.path.join(
args.save_dir,
'seqsupce_{}+{}_cpb{}_epo{}_bs{}_lr{}_mom{}_wd{}_lcr{}{}'.format(
args.arch, args.data, args.cpb, args.epo, args.bs, args.lr, args.mom,
args.wd, args.lcr, args.tail))
headers = ["label_batch_id", "epoch", "step", "loss"]
utils.init_pipeline(model_dir, headers)
utils.save_params(model_dir, vars(args))
## per model functions
def lr_schedule(epoch, optimizer):
"""decrease the learning rate"""
lr = args.lr
if epoch >= 400:
lr = args.lr * 0.01
elif epoch >= 200:
lr = args.lr * 0.1
for param_group in optimizer.param_groups:
param_group['lr'] = lr
help='base directory for saving PyTorch model. (default: ./saved_models/)')
parser.add_argument(
'--data_dir',
type=str,
default='./data/',
help='base directory for saving PyTorch model. (default: ./data/)')
args = parser.parse_args()
## Pipelines Setup
model_dir = os.path.join(
args.save_dir,
'selfsup_{}+{}_{}_epo{}_bs{}_aug{}+{}_lr{}_mom{}_wd{}_gam1{}_gam2{}_eps{}{}'
.format(args.arch, args.fd, args.data, args.epo, args.bs, args.aug,
args.transform, args.lr, args.mom, args.wd, args.gam1, args.gam2,
args.eps, args.tail))
utils.init_pipeline(model_dir)
## Prepare for Training
if args.pretrain_dir is not None:
net, _ = tf.load_checkpoint(args.pretrain_dir, args.pretrain_epo)
utils.update_params(model_dir, args.pretrain_dir)
else:
net = tf.load_architectures(args.arch, args.fd)
transforms = tf.load_transforms(args.transform)
trainset = tf.load_trainset(args.data, path=args.data_dir)
trainloader = AugmentLoader(trainset,
transforms=transforms,
sampler=args.sampler,
batch_size=args.bs,
num_aug=args.aug)