def main():
global best_acc1
best_acc1 = 0
args = parse_option()
if args.gpu is not None:
print("Use GPU: {} for training".format(args.gpu))
# set the data loader
train_folder = os.path.join(args.data_folder, 'train')
val_folder = os.path.join(args.data_folder, 'val')
logger = getLogger(args.save_folder)
if args.dataset.startswith('imagenet') or args.dataset.startswith(
'places'):
image_size = 224
crop_padding = 32
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
normalize = transforms.Normalize(mean=mean, std=std)
if args.aug == 'NULL':
train_transform = transforms.Compose([
transforms.RandomResizedCrop(image_size,
scale=(args.crop, 1.)),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
])
elif args.aug == 'CJ':
train_transform = transforms.Compose([
transforms.RandomResizedCrop(image_size,
scale=(args.crop, 1.)),
transforms.RandomGrayscale(p=0.2),
transforms.ColorJitter(0.4, 0.4, 0.4, 0.4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
])
else:
raise NotImplemented('augmentation not supported: {}'.format(
args.aug))
val_transform = transforms.Compose([
transforms.Resize(image_size + crop_padding),
transforms.CenterCrop(image_size),
transforms.ToTensor(),
normalize,
])
if args.dataset.startswith('imagenet'):
train_dataset = datasets.ImageFolder(train_folder, train_transform)
val_dataset = datasets.ImageFolder(
val_folder,
val_transform,
)
if args.dataset.startswith('places'):
train_dataset = ImageList(
'/data/trainvalsplit_places205/train_places205.csv',
'/data/data/vision/torralba/deeplearning/images256',
transform=train_transform,
symbol_split=' ')
val_dataset = ImageList(
'/data/trainvalsplit_places205/val_places205.csv',
'/data/data/vision/torralba/deeplearning/images256',
transform=val_transform,
symbol_split=' ')
print(len(train_dataset))
train_sampler = None
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=args.batch_size,
shuffle=(train_sampler is None),
num_workers=args.n_workers,
pin_memory=False,
sampler=train_sampler)
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.n_workers,
pin_memory=False)
elif args.dataset.startswith('cifar'):
train_loader, val_loader = cifar.get_linear_dataloader(args)
elif args.dataset.startswith('svhn'):
train_loader, val_loader = svhn.get_linear_dataloader(args)
# create model and optimizer
if args.model == 'alexnet':
if args.layer == 6:
args.layer = 5
model = AlexNet(128)
model = nn.DataParallel(model)
classifier = LinearClassifierAlexNet(args.layer, args.n_label, 'avg')
elif args.model == 'alexnet_cifar':
if args.layer == 6:
args.layer = 5
model = AlexNet_cifar(128)
model = nn.DataParallel(model)
classifier = LinearClassifierAlexNet(args.layer,
args.n_label,
'avg',
cifar=True)
elif args.model == 'resnet50':
model = resnet50(non_linear_head=False)
model = nn.DataParallel(model)
classifier = LinearClassifierResNet(args.layer, args.n_label, 'avg', 1)
elif args.model == 'resnet18':
model = resnet18()
model = nn.DataParallel(model)
classifier = LinearClassifierResNet(args.layer,
args.n_label,
'avg',
1,
bottleneck=False)
elif args.model == 'resnet18_cifar':
model = resnet18_cifar()
model = nn.DataParallel(model)
classifier = LinearClassifierResNet(args.layer,
args.n_label,
'avg',
1,
bottleneck=False)
elif args.model == 'resnet50_cifar':
model = resnet50_cifar()
model = nn.DataParallel(model)
classifier = LinearClassifierResNet(args.layer, args.n_label, 'avg', 1)
elif args.model == 'resnet50x2':
model = InsResNet50(width=2)
classifier = LinearClassifierResNet(args.layer, args.n_label, 'avg', 2)
elif args.model == 'resnet50x4':
model = InsResNet50(width=4)
classifier = LinearClassifierResNet(args.layer, args.n_label, 'avg', 4)
elif args.model == 'shufflenet':
model = shufflenet_v2_x1_0(num_classes=128, non_linear_head=False)
model = nn.DataParallel(model)
classifier = LinearClassifierResNet(args.layer, args.n_label, 'avg',
0.5)
else:
raise NotImplementedError('model not supported {}'.format(args.model))
print('==> loading pre-trained model')
ckpt = torch.load(args.model_path)
if not args.moco:
model.load_state_dict(ckpt['state_dict'])
else:
try:
state_dict = ckpt['state_dict']
for k in list(state_dict.keys()):
# retain only encoder_q up to before the embedding layer
if k.startswith('module.encoder_q'
) and not k.startswith('module.encoder_q.fc'):
# remove prefix
state_dict['module.' +
k[len("module.encoder_q."):]] = state_dict[k]
# delete renamed or unused k
del state_dict[k]
model.load_state_dict(state_dict)
except:
pass
print("==> loaded checkpoint '{}' (epoch {})".format(
args.model_path, ckpt['epoch']))
print('==> done')
model = model.cuda()
classifier = classifier.cuda()
criterion = torch.nn.CrossEntropyLoss().cuda(args.gpu)
if not args.adam:
optimizer = torch.optim.SGD(classifier.parameters(),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
else:
optimizer = torch.optim.Adam(classifier.parameters(),
lr=args.learning_rate,
betas=(args.beta1, args.beta2),
weight_decay=args.weight_decay,
eps=1e-8)
model.eval()
cudnn.benchmark = True
# optionally resume from a checkpoint
args.start_epoch = 1
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume, map_location='cpu')
# checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch'] + 1
classifier.load_state_dict(checkpoint['classifier'])
optimizer.load_state_dict(checkpoint['optimizer'])
best_acc1 = checkpoint['best_acc1']
print(best_acc1.item())
best_acc1 = best_acc1.cuda()
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
if 'opt' in checkpoint.keys():
# resume optimization hyper-parameters
print('=> resume hyper parameters')
if 'bn' in vars(checkpoint['opt']):
print('using bn: ', checkpoint['opt'].bn)
if 'adam' in vars(checkpoint['opt']):
print('using adam: ', checkpoint['opt'].adam)
#args.learning_rate = checkpoint['opt'].learning_rate
# args.lr_decay_epochs = checkpoint['opt'].lr_decay_epochs
args.lr_decay_rate = checkpoint['opt'].lr_decay_rate
args.momentum = checkpoint['opt'].momentum
args.weight_decay = checkpoint['opt'].weight_decay
args.beta1 = checkpoint['opt'].beta1
args.beta2 = checkpoint['opt'].beta2
del checkpoint
torch.cuda.empty_cache()
else:
print("=> no checkpoint found at '{}'".format(args.resume))
# tensorboard
tblogger = tb_logger.Logger(logdir=args.tb_folder, flush_secs=2)
# routine
best_acc = 0.0
for epoch in range(args.start_epoch, args.epochs + 1):
adjust_learning_rate(epoch, args, optimizer)
print("==> training...")
time1 = time.time()
train_acc, train_acc5, train_loss = train(epoch, train_loader, model,
classifier, criterion,
optimizer, args)
time2 = time.time()
logging.info('train epoch {}, total time {:.2f}'.format(
epoch, time2 - time1))
logging.info(
'Epoch: {}, lr:{} , train_loss: {:.4f}, train_acc: {:.4f}/{:.4f}'.
format(epoch, optimizer.param_groups[0]['lr'], train_loss,
train_acc, train_acc5))
tblogger.log_value('train_acc', train_acc, epoch)
tblogger.log_value('train_acc5', train_acc5, epoch)
tblogger.log_value('train_loss', train_loss, epoch)
tblogger.log_value('learning_rate', optimizer.param_groups[0]['lr'],
epoch)
test_acc, test_acc5, test_loss = validate(val_loader, model,
classifier, criterion, args)
if test_acc >= best_acc:
best_acc = test_acc
logging.info(
colorful(
'Epoch: {}, val_loss: {:.4f}, val_acc: {:.4f}/{:.4f}, best_acc: {:.4f}'
.format(epoch, test_loss, test_acc, test_acc5, best_acc)))
tblogger.log_value('test_acc', test_acc, epoch)
tblogger.log_value('test_acc5', test_acc5, epoch)
tblogger.log_value('test_loss', test_loss, epoch)
# save the best model
if test_acc > best_acc1:
best_acc1 = test_acc
state = {
'opt': args,
'epoch': epoch,
'classifier': classifier.state_dict(),
'best_acc1': best_acc1,
'optimizer': optimizer.state_dict(),
}
save_name = '{}_layer{}.pth'.format(args.model, args.layer)
save_name = os.path.join(args.save_folder, save_name)
print('saving best model!')
torch.save(state, save_name)
# save model
if epoch % args.save_freq == 0:
print('==> Saving...')
state = {
'opt': args,
'epoch': epoch,
'classifier': classifier.state_dict(),
'best_acc1': test_acc,
'optimizer': optimizer.state_dict(),
}
save_name = 'ckpt_epoch_{epoch}.pth'.format(epoch=epoch)
save_name = os.path.join(args.save_folder, save_name)
print('saving regular model!')
torch.save(state, save_name)
# tensorboard logger
pass
def extract_feature(args):
"""Extract and save features for train split, several clips per video."""
torch.backends.cudnn.benchmark = True
# Force the pytorch to create context on the specific device
os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu)
device = torch.device('cuda:0' if torch.cuda.is_available() else "cpu")
########### model ##############
model = AlexNet(with_classifier=False, return_conv=True) .to(device)
if args.ckpt:
pretrained_weights = load_pretrained_weights(args.ckpt)
model.load_state_dict(pretrained_weights, strict=True)
model.eval()
torch.set_grad_enabled(False)
### Exract for train split ###
train_transforms = transforms.Compose([
transforms.Resize((256, 256)),
transforms.CenterCrop(224),
transforms.ToTensor()
])
train_dataset = UCF101FrameRetrievalDataset('data/ucf101', 10, True, train_transforms)
train_dataloader = DataLoader(train_dataset, batch_size=args.bs, shuffle=False,
num_workers=args.workers, pin_memory=True, drop_last=True)
features = []
classes = []
for data in tqdm(train_dataloader):
sampled_clips, idxs = data
clips = sampled_clips.reshape((-1, 3, 224, 224))
inputs = clips.to(device)
# forward
outputs = model(inputs)
# print(outputs.shape)
# exit()
features.append(outputs.cpu().numpy().tolist())
classes.append(idxs.cpu().numpy().tolist())
features = np.array(features).reshape(-1, 10, outputs.shape[1])
classes = np.array(classes).reshape(-1, 10)
np.save(os.path.join(args.feature_dir, 'train_feature.npy'), features)
np.save(os.path.join(args.feature_dir, 'train_class.npy'), classes)
### Exract for test split ###
test_transforms = transforms.Compose([
transforms.Resize((256, 256)),
transforms.CenterCrop(224),
transforms.ToTensor()
])
test_dataset = UCF101FrameRetrievalDataset('data/ucf101', 10, False, test_transforms)
test_dataloader = DataLoader(test_dataset, batch_size=args.bs, shuffle=False,
num_workers=args.workers, pin_memory=True, drop_last=True)
features = []
classes = []
for data in tqdm(test_dataloader):
sampled_clips, idxs = data
clips = sampled_clips.reshape((-1, 3, 224, 224))
inputs = clips.to(device)
# forward
outputs = model(inputs)
features.append(outputs.cpu().numpy().tolist())
classes.append(idxs.cpu().numpy().tolist())
features = np.array(features).reshape(-1, 10, outputs.shape[1])
classes = np.array(classes).reshape(-1, 10)
np.save(os.path.join(args.feature_dir, 'test_feature.npy'), features)
np.save(os.path.join(args.feature_dir, 'test_class.npy'), classes)
pretrain_model_path = './AlexNet/alexnet-owt-4df8aa71.pth'
n_validate_batch = 100 # Number of batches used for validation
validate_batch_size = 50 # Batch size of validation
prune_ratio = {
'features.0': 80,
'features.3': 35,
'features.6': 35,
'features.8': 35,
'features.10': 35,
'classifier.1': 10,
'classifier.4': 10,
'classifier.6': 35
}
# -------------------------------------------- User Config ------------------------------------
net = AlexNet()
net.load_state_dict(torch.load(pretrain_model_path))
param = net.state_dict()
total_nnz = 0
total_nelements = 0
for layer_name, CR in prune_ratio.items():
if CR != 100:
pruned_weight = np.load('%s/CR_%d/%s.weight.npy' %
(pruned_weight_root, CR, layer_name))
pruned_bias = np.load('%s/CR_%d/%s.bias.npy' %
(pruned_weight_root, CR, layer_name))
# Calculate sparsity
total_nnz += np.count_nonzero(pruned_weight)
total_nnz += np.count_nonzero(pruned_bias)
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--stage', default='train', type=str)
parser.add_argument('--dataset', default='imagenet', type=str)
parser.add_argument('--lr', default=0.0012, type=float)
parser.add_argument('--batch_size', default=128, type=int)
parser.add_argument('--gpus', default='0,1,2,3', type=str)
parser.add_argument('--weight_decay', default=1e-5, type=float)
parser.add_argument('--max_epoch', default=30, type=int)
parser.add_argument('--lr_decay_steps', default='15,20,25', type=str)
parser.add_argument('--exp', default='', type=str)
parser.add_argument('--list', default='', type=str)
parser.add_argument('--resume_path', default='', type=str)
parser.add_argument('--pretrain_path', default='', type=str)
parser.add_argument('--n_workers', default=32, type=int)
parser.add_argument('--network', default='resnet50', type=str)
global args
args = parser.parse_args()
if not os.path.exists(args.exp):
os.makedirs(args.exp)
if not os.path.exists(os.path.join(args.exp, 'runs')):
os.makedirs(os.path.join(args.exp, 'runs'))
if not os.path.exists(os.path.join(args.exp, 'models')):
os.makedirs(os.path.join(args.exp, 'models'))
if not os.path.exists(os.path.join(args.exp, 'logs')):
os.makedirs(os.path.join(args.exp, 'logs'))
# logger initialize
logger = getLogger(args.exp)
device_ids = list(map(lambda x: int(x), args.gpus.split(',')))
device = torch.device('cuda: 0')
train_loader, val_loader = cifar.get_semi_dataloader(
args) if args.dataset.startswith(
'cifar') else imagenet.get_semi_dataloader(args)
# create model
if args.network == 'alexnet':
network = AlexNet(128)
elif args.network == 'alexnet_cifar':
network = AlexNet_cifar(128)
elif args.network == 'resnet18_cifar':
network = ResNet18_cifar()
elif args.network == 'resnet50_cifar':
network = ResNet50_cifar()
elif args.network == 'wide_resnet28':
network = WideResNet(28, args.dataset == 'cifar10' and 10 or 100, 2)
elif args.network == 'resnet18':
network = resnet18()
elif args.network == 'resnet50':
network = resnet50()
network = nn.DataParallel(network, device_ids=device_ids)
network.to(device)
classifier = nn.Linear(2048, 1000).to(device)
# create optimizer
parameters = network.parameters()
optimizer = torch.optim.SGD(
parameters,
lr=args.lr,
momentum=0.9,
weight_decay=args.weight_decay,
)
cls_optimizer = torch.optim.SGD(
classifier.parameters(),
lr=args.lr * 50,
momentum=0.9,
weight_decay=args.weight_decay,
)
cudnn.benchmark = True
# create memory_bank
global writer
writer = SummaryWriter(comment='SemiSupervised',
logdir=os.path.join(args.exp, 'runs'))
# create criterion
criterion = nn.CrossEntropyLoss()
logging.info(beautify(args))
start_epoch = 0
if args.pretrain_path != '' and args.pretrain_path != 'none':
logging.info('loading pretrained file from {}'.format(
args.pretrain_path))
checkpoint = torch.load(args.pretrain_path)
state_dict = checkpoint['state_dict']
valid_state_dict = {
k: v
for k, v in state_dict.items()
if k in network.state_dict() and 'fc.' not in k
}
for k, v in network.state_dict().items():
if k not in valid_state_dict:
logging.info('{}: Random Init'.format(k))
valid_state_dict[k] = v
# logging.info(valid_state_dict.keys())
network.load_state_dict(valid_state_dict)
else:
logging.info('Training SemiSupervised Learning From Scratch')
logging.info('start training')
best_acc = 0.0
try:
for i_epoch in range(start_epoch, args.max_epoch):
train(i_epoch, network, classifier, criterion, optimizer,
cls_optimizer, train_loader, device)
checkpoint = {
'epoch': i_epoch + 1,
'state_dict': network.state_dict(),
'optimizer': optimizer.state_dict(),
}
torch.save(checkpoint,
os.path.join(args.exp, 'models', 'checkpoint.pth'))
adjust_learning_rate(args.lr_decay_steps, optimizer, i_epoch)
if i_epoch % 2 == 0:
acc1, acc5 = validate(i_epoch, network, classifier, val_loader,
device)
if acc1 >= best_acc:
best_acc = acc1
torch.save(checkpoint,
os.path.join(args.exp, 'models', 'best.pth'))
writer.add_scalar('acc1', acc1, i_epoch + 1)
writer.add_scalar('acc5', acc5, i_epoch + 1)
if i_epoch in [30, 60, 120, 160, 200]:
torch.save(
checkpoint,
os.path.join(args.exp, 'models',
'{}.pth'.format(i_epoch + 1)))
logging.info(
colorful('[Epoch: {}] val acc: {:.4f}/{:.4f}'.format(
i_epoch, acc1, acc5)))
logging.info(
colorful('[Epoch: {}] best acc: {:.4f}'.format(
i_epoch, best_acc)))
with torch.no_grad():
for name, param in network.named_parameters():
if 'bn' not in name:
writer.add_histogram(name, param, i_epoch)
# cluster
except KeyboardInterrupt as e:
logging.info('KeyboardInterrupt at {} Epochs'.format(i_epoch))
exit()
def main():
# Init logger
if not os.path.isdir(args.save_path):
os.makedirs(args.save_path)
log = open(os.path.join(args.save_path, 'log_seed_{}.txt'.format(args.manualSeed)), 'w')
print_log('save path : {}'.format(args.save_path), log)
state = {k: v for k, v in args._get_kwargs()}
print_log(state, log)
print_log("Random Seed: {}".format(args.manualSeed), log)
print_log("python version : {}".format(sys.version.replace('\n', ' ')), log)
print_log("torch version : {}".format(torch.__version__), log)
print_log("cudnn version : {}".format(torch.backends.cudnn.version()), log)
# Data loading code
# Any other preprocessings? http://pytorch.org/audio/transforms.html
sample_length = 10000
scale = transforms.Scale()
padtrim = transforms.PadTrim(sample_length)
downmix = transforms.DownmixMono()
transforms_audio = transforms.Compose([
scale, padtrim, downmix
])
if not os.path.isdir(args.data_path):
os.makedirs(args.data_path)
train_dir = os.path.join(args.data_path, 'train')
val_dir = os.path.join(args.data_path, 'val')
#Choose dataset to use
if args.dataset == 'arctic':
# TODO No ImageFolder equivalent for audio. Need to create a Dataset manually
train_dataset = Arctic(train_dir, transform=transforms_audio, download=True)
val_dataset = Arctic(val_dir, transform=transforms_audio, download=True)
num_classes = 4
elif args.dataset == 'vctk':
train_dataset = dset.VCTK(train_dir, transform=transforms_audio, download=True)
val_dataset = dset.VCTK(val_dir, transform=transforms_audio, download=True)
num_classes = 10
elif args.dataset == 'yesno':
train_dataset = dset.YESNO(train_dir, transform=transforms_audio, download=True)
val_dataset = dset.YESNO(val_dir, transform=transforms_audio, download=True)
num_classes = 2
else:
assert False, 'Dataset is incorrect'
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
# pin_memory=True, # What is this?
# sampler=None # What is this?
)
val_loader = torch.utils.data.DataLoader(
val_dataset,
batch_size=args.batch_size, shuffle=False,
num_workers=args.workers, pin_memory=True)
#Feed in respective model file to pass into model (alexnet.py)
print_log("=> creating model '{}'".format(args.arch), log)
# Init model, criterion, and optimizer
# net = models.__dict__[args.arch](num_classes)
net = AlexNet(num_classes)
#
print_log("=> network :\n {}".format(net), log)
# net = torch.nn.DataParallel(net, device_ids=list(range(args.ngpu)))
# define loss function (criterion) and optimizer
criterion = torch.nn.CrossEntropyLoss()
# Define stochastic gradient descent as optimizer (run backprop on random small batch)
optimizer = torch.optim.SGD(net.parameters(), state['learning_rate'], momentum=state['momentum'],
weight_decay=state['decay'], nesterov=True)
#Sets use for GPU if available
if args.use_cuda:
net.cuda()
criterion.cuda()
recorder = RecorderMeter(args.epochs)
# optionally resume from a checkpoint
# Need same python vresion that the resume was in
if args.resume:
if os.path.isfile(args.resume):
print_log("=> loading checkpoint '{}'".format(args.resume), log)
if args.ngpu == 0:
checkpoint = torch.load(args.resume, map_location=lambda storage, loc: storage)
else:
checkpoint = torch.load(args.resume)
recorder = checkpoint['recorder']
args.start_epoch = checkpoint['epoch']
net.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print_log("=> loaded checkpoint '{}' (epoch {})" .format(args.resume, checkpoint['epoch']), log)
else:
print_log("=> no checkpoint found at '{}'".format(args.resume), log)
else:
print_log("=> do not use any checkpoint for {} model".format(args.arch), log)
if args.evaluate:
validate(val_loader, net, criterion, 0, log, val_dataset)
return
# Main loop
start_time = time.time()
epoch_time = AverageMeter()
# Training occurs here
for epoch in range(args.start_epoch, args.epochs):
current_learning_rate = adjust_learning_rate(optimizer, epoch, args.gammas, args.schedule)
need_hour, need_mins, need_secs = convert_secs2time(epoch_time.avg * (args.epochs-epoch))
need_time = '[Need: {:02d}:{:02d}:{:02d}]'.format(need_hour, need_mins, need_secs)
print_log('\n==>>{:s} [Epoch={:03d}/{:03d}] {:s} [learning_rate={:6.4f}]'.format(time_string(), epoch, args.epochs, need_time, current_learning_rate) \
+ ' [Best : Accuracy={:.2f}, Error={:.2f}]'.format(recorder.max_accuracy(False), 100-recorder.max_accuracy(False)), log)
print("One epoch")
# train for one epoch
# Call to train (note that our previous net is passed into the model argument)
train_acc, train_los = train(train_loader, net, criterion, optimizer, epoch, log, train_dataset)
# evaluate on validation set
#val_acc, val_los = extract_features(test_loader, net, criterion, log)
val_acc, val_los = validate(val_loader, net, criterion, epoch, log, val_dataset)
is_best = recorder.update(epoch, train_los, train_acc, val_los, val_acc)
save_checkpoint({
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': net.state_dict(),
'recorder': recorder,
'optimizer' : optimizer.state_dict(),
}, is_best, args.save_path, 'checkpoint.pth.tar')
# measure elapsed time
epoch_time.update(time.time() - start_time)
start_time = time.time()
recorder.plot_curve( os.path.join(args.save_path, 'curve.png') )
log.close()
P += (pred_cls == input_label_tensor).sum().cpu().detach().numpy()
N += HyperParams["batch_size"]
if idx % 500 == 499:
print("|acc:%f|use time:%s|" %
(float(P / N), str(time.time() - start_time)))
start_time = time.time()
# print('')
if __name__ == '__main__':
train_data = mnist.MNIST("./mnist_data")
model = AlexNet(10)
if HyperParams["cuda"]:
model = model.cuda()
optimer = torch.optim.Adam(params=[{
"params": model.parameters()
}],
lr=0.004)
lr_sch = torch.optim.lr_scheduler.MultiStepLR(optimer, [1, 2, 3, 4], 0.1)
criterion = torch.nn.CrossEntropyLoss()
static_params = torch.load("./%s_E%d.snap" %
(HyperParams["model_save_prefix"], 4))
model.load_state_dict(static_params)
# trainval(model,optimer,lr_sch,criterion,train_data)
if HyperParams["quantize"]:
model = torch.quantization.quantize_dynamic(model)
torch.save(model.state_dict(), "./quantize_mode.snap")
test(model, train_data)
def evaluate():
num_classes = 4
# Init logger
if not os.path.isdir(args.save_path):
os.makedirs(args.save_path)
log = open(os.path.join(args.save_path, 'log_seed_{}.txt'.format(args.manualSeed)), 'w')
print_log('save path : {}'.format(args.save_path), log)
state = {k: v for k, v in args._get_kwargs()}
print_log(state, log)
print_log("Random Seed: {}".format(args.manualSeed), log)
print_log("python version : {}".format(sys.version.replace('\n', ' ')), log)
print_log("torch version : {}".format(torch.__version__), log)
print_log("cudnn version : {}".format(torch.backends.cudnn.version()), log)
# Any other preprocessings? http://pytorch.org/audio/transforms.html
sample_length = 10000
scale = transforms.Scale()
padtrim = transforms.PadTrim(sample_length)
transforms_audio = transforms.Compose([
scale, padtrim
])
# Data loading
fs, data = wavfile.read(args.file_name)
data = torch.from_numpy(data).float()
data = data.unsqueeze(1)
audio = transforms_audio(data)
audio = Variable(audio)
audio = audio.view(1, -1)
audio = audio.unsqueeze(0)
#Feed in respective model file to pass into model (alexnet.py)
print_log("=> creating model '{}'".format(args.arch), log)
# Init model, criterion, and optimizer
# net = models.__dict__[args.arch](num_classes)
net = AlexNet(num_classes)
print_log("=> network :\n {}".format(net), log)
#Sets use for GPU if available
if args.use_cuda:
net.cuda()
# optionally resume from a checkpoint
# Need same python version that the resume was in
if args.resume:
if os.path.isfile(args.resume):
print_log("=> loading checkpoint '{}'".format(args.resume), log)
if args.ngpu == 0:
checkpoint = torch.load(args.resume, map_location=lambda storage, loc: storage)
else:
checkpoint = torch.load(args.resume)
recorder = checkpoint['recorder']
args.start_epoch = checkpoint['epoch']
net.load_state_dict(checkpoint['state_dict'])
print_log("=> loaded checkpoint '{}' (epoch {})" .format(args.resume, checkpoint['epoch']), log)
else:
print_log("=> no checkpoint found at '{}'".format(args.resume), log)
else:
print_log("=> do not use any checkpoint for {} model".format(args.arch), log)
net.eval()
if args.use_cuda:
audio = audio.cuda()
output = net(audio)
print(output)
# TODO postprocess output to a string representing the person speaking
# ouptut = val_dataset.postprocess_target(output)
return