def main():
# [b, 32, 32, 3] => [b, 1, 1, 512]
model = resnet10()
model.build(input_shape=(None, 32, 32, 3)) # 这里input_shape 用 [] 就会报错，不知道为啥
model.summary()
optimizer = optimizers.Adam(1e-4)
for epoch in range(50):
for step, (x, y) in enumerate(train_db):
with tf.GradientTape() as tape:
# [b, 32, 32, 3] => [b, 100]
logits = model(x, training=True)
y_onehot = tf.one_hot(y, depth=100)
loss = tf.losses.categorical_crossentropy(y_onehot, logits, from_logits=True)
loss = tf.reduce_mean(loss)
grads = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(grads, model.trainable_variables))
if step%100 == 0:
print(epoch, step, 'loss:', float(loss))
total_num = 0
total_correct = 0
for x, y in test_db:
logits = model(x, training=False)
prob = tf.nn.softmax(logits, axis=1)
pred = tf.argmax(prob, axis=1)
pred = tf.cast(pred, dtype=tf.int32)
correct = tf.cast(tf.equal(pred, y), dtype=tf.int32)
correct = tf.reduce_sum(correct)
total_num += x.shape[0]
total_correct += int(correct)
acc = total_correct / total_num
print(epoch, 'acc:', acc)
def generate_model(sample_input_W,
sample_input_H,
sample_input_D,
num_seg_classes=3,
no_cuda=False,
phase='train',
pretrain_path=None,
new_layer_names=['avgpool', 'fc']):
model = resnet.resnet10(sample_input_W=sample_input_W,
sample_input_H=sample_input_H,
sample_input_D=sample_input_D,
num_seg_classes=num_seg_classes)
if not no_cuda:
if torch.cuda.device_count() > 1:
model = model.cuda()
model = nn.DataParallel(model,
device_ids=range(
torch.cuda.device_count()))
net_dict = model.state_dict()
else:
import os
os.environ["CUDA_VISIBLE_DEVICES"] = str(0)
model = model.cuda()
model = nn.DataParallel(model, device_ids=None)
net_dict = model.state_dict()
else:
net_dict = model.state_dict()
# load pretrain
if phase != 'test' and pretrain_path:
print('loading pretrained model {}'.format(pretrain_path))
pretrain = torch.load(pretrain_path)
pretrain_dict = {
k: v
for k, v in pretrain['state_dict'].items() if k in net_dict.keys()
}
net_dict.update(pretrain_dict)
model.load_state_dict(net_dict)
new_parameters = []
for pname, p in model.named_parameters():
for layer_name in new_layer_names:
if pname.find(layer_name) >= 0:
new_parameters.append(p)
break
new_parameters_id = list(map(id, new_parameters))
base_parameters = list(
filter(lambda p: id(p) not in new_parameters_id,
model.parameters()))
parameters = {
'base_parameters': base_parameters,
'new_parameters': new_parameters
}
return model, parameters
return model, model.parameters()
import time
from load_data_v3 import train_dataloader, test_dataloader
import json
# from c3d_v1 import C3D
import resnet
# import resnet_v2
# import deepmedic
# import vggnet_v1
config = json.load(open('package.json'))
num_epoch = config['epoch']
os.environ['CUDA_VISIBLE_DEVICES'] = '0' # 第3块GPU
DEVICE = torch.device('cuda:0')
# c3d = C3D()
c3d = resnet.resnet10(sample_size=8, sample_duration=4)
# c3d = resnet_v2.generate_model(model_depth=10)
# c3d = deepmedic.UNet3D(in_channel=2,n_classes=2)
# c3d = vggnet_v1.vgg13_bn()
c3d.to(DEVICE)
# 优化器和损失函数\
lr = 0.01
optimizer = torch.optim.SGD(c3d.parameters(), lr=lr, weight_decay=0.01)
# optimizer = torch.optim.Adam(c3d.parameters(),lr=lr,betas=(0.9,0.999))
criterion = torch.nn.CrossEntropyLoss()
# criterion = torch.nn.BCELoss(reduce=False, size_average=False)
tmp = []
testtmp = []
loss1 = []
def generate_model(opt):
assert opt.model in [
'resnet', 'preresnet', 'wideresnet', 'resnext', 'densenet'
]
if opt.model == 'resnet':
assert opt.model_depth in [10, 18, 34, 50, 101, 152, 200]
from resnet import get_fine_tuning_parameters
if opt.model_depth == 10:
model = resnet.resnet10(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 18:
model = resnet.resnet18(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 34:
model = resnet.resnet34(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 50:
model = resnet.resnet50(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 101:
model = resnet.resnet101(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 152:
model = resnet.resnet152(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 200:
model = resnet.resnet200(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model == 'wideresnet':
assert opt.model_depth in [50]
from models.wide_resnet import get_fine_tuning_parameters
if opt.model_depth == 50:
model = wide_resnet.resnet50(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
k=opt.wide_resnet_k,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model == 'resnext':
assert opt.model_depth in [50, 101, 152]
from models.resnext import get_fine_tuning_parameters
if opt.model_depth == 50:
model = resnext.resnet50(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
cardinality=opt.resnext_cardinality,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 101:
model = resnext.resnet101(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
cardinality=opt.resnext_cardinality,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 152:
model = resnext.resnet152(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
cardinality=opt.resnext_cardinality,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model == 'preresnet':
assert opt.model_depth in [18, 34, 50, 101, 152, 200]
from models.pre_act_resnet import get_fine_tuning_parameters
if opt.model_depth == 18:
model = pre_act_resnet.resnet18(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 34:
model = pre_act_resnet.resnet34(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 50:
model = pre_act_resnet.resnet50(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 101:
model = pre_act_resnet.resnet101(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 152:
model = pre_act_resnet.resnet152(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 200:
model = pre_act_resnet.resnet200(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model == 'densenet':
assert opt.model_depth in [121, 169, 201, 264]
from models.densenet import get_fine_tuning_parameters
if opt.model_depth == 121:
model = densenet.densenet121(num_classes=opt.n_classes,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 169:
model = densenet.densenet169(num_classes=opt.n_classes,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 201:
model = densenet.densenet201(num_classes=opt.n_classes,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 264:
model = densenet.densenet264(num_classes=opt.n_classes,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
if not opt.no_cuda:
model = model.cuda()
model = nn.DataParallel(model, device_ids=None)
if opt.pretrain_path:
print('loading pretrained model {}'.format(opt.pretrain_path))
pretrain = torch.load(opt.pretrain_path)
assert opt.arch == pretrain['arch']
model.load_state_dict(pretrain['state_dict'])
if opt.model == 'densenet':
model.module.classifier = nn.Linear(
model.module.classifier.in_features,
opt.n_finetune_classes)
model.module.classifier = model.module.classifier.cuda()
else:
model.module.fc = nn.Linear(model.module.fc.in_features,
opt.n_finetune_classes)
model.module.fc = model.module.fc.cuda()
parameters = get_fine_tuning_parameters(model, opt.ft_begin_index)
return model, parameters
else:
if opt.pretrain_path:
print('loading pretrained model {}'.format(opt.pretrain_path))
pretrain = torch.load(opt.pretrain_path)
assert opt.arch == pretrain['arch']
model.load_state_dict(pretrain['state_dict'])
if opt.model == 'densenet':
model.classifier = nn.Linear(model.classifier.in_features,
opt.n_finetune_classes)
else:
model.fc = nn.Linear(model.fc.in_features,
opt.n_finetune_classes)
parameters = get_fine_tuning_parameters(model, opt.ft_begin_index)
return model, parameters
return model, model.parameters()
def main(args):
args.cuda = not args.no_cuda and torch.cuda.is_available()
init_seeds(seed=int(time.time()))
kwargs = {'num_workers': 2, 'pin_memory': True} if args.cuda else {}
print(args.dataset)
if args.dataset == 'MNIST':
test_dataloader = data.DataLoader(
MNIST(args.data_path, args.run_folder, transform=mnist_transformer()),
batch_size=10000, shuffle=False, **kwargs)
train_dataset = MNIST(args.data_path, args.run_folder, train=True, transform=mnist_transformer(), imbalance_ratio=args.imbalance_ratio)
if args.imbalance_ratio == 100:
args.num_images = 25711
else:
args.num_images = 50000
args.budget = 125
args.initial_budget = 125
args.num_classes = 10
args.num_channels = 1
args.arch_scaler = 2
elif args.dataset == 'SVHN':
test_dataloader = data.DataLoader(
SVHN(args.data_path, args.run_folder, transform=svhn_transformer()),
batch_size=5000, shuffle=False, **kwargs)
train_dataset = SVHN(args.data_path, args.run_folder, train=True, transform=svhn_transformer(), imbalance_ratio=args.imbalance_ratio)
if args.imbalance_ratio == 100:
args.num_images = 318556
else:
args.num_images = 500000
args.budget = 1250
args.initial_budget = 1250
args.num_classes = 10
args.num_channels = 3
args.arch_scaler = 1
elif args.dataset == 'cifar10':
test_dataloader = data.DataLoader(
datasets.CIFAR10(args.data_path, download=True, transform=cifar_transformer(), train=False),
batch_size=args.batch_size, drop_last=False)
train_dataset = CIFAR10(args.data_path)
args.num_images = 50000
args.budget = 2500
args.initial_budget = 5000
args.num_classes = 10
args.num_channels = 3
elif args.dataset == 'cifar100':
test_dataloader = data.DataLoader(
datasets.CIFAR100(args.data_path, download=True, transform=cifar_transformer(), train=False),
batch_size=args.batch_size, drop_last=False)
train_dataset = CIFAR100(args.data_path)
args.num_images = 50000
args.budget = 2500
args.initial_budget = 5000
args.num_classes = 100
args.num_channels = 3
elif args.dataset == 'ImageNet':
test_dataloader = data.DataLoader(
ImageNet(args.data_path + '/val', transform=imagenet_test_transformer()),
batch_size=args.batch_size, shuffle=False, drop_last=False, **kwargs)
if args.imbalance_ratio == 100:
train_dataset = ImageNet(args.data_path + '/train_ir_100', transform=imagenet_train_transformer())
args.num_images = 645770
else:
train_dataset = ImageNet(args.data_path + '/train', transform=imagenet_train_transformer())
args.num_images = 1281167
args.budget = 64000
args.initial_budget = 64000
args.num_classes = 1000
args.num_channels = 3
args.arch_scaler = 1
else:
raise NotImplementedError
all_indices = set(np.arange(args.num_images))
initial_indices = random.sample(all_indices, args.initial_budget)
sampler = data.sampler.SubsetRandomSampler(initial_indices)
#print(args.batch_size, sampler)
# dataset with labels available
querry_dataloader = data.DataLoader(train_dataset, sampler=sampler,
batch_size=args.batch_size, drop_last=False, **kwargs)
print('Sampler size =', len(querry_dataloader))
solver = Solver(args, test_dataloader)
splits = range(1,11)
current_indices = list(initial_indices)
accuracies = []
for split in splits:
print("Split =", split)
# need to retrain all the models on the new images
# re initialize and retrain the models
#task_model = vgg.vgg16_bn(num_classes=args.num_classes)
if args.dataset == 'MNIST':
task_model = model.LeNet(num_classes=args.num_classes)
elif args.dataset == 'SVHN':
task_model = resnet.resnet10(num_classes=args.num_classes)
elif args.dataset == 'ImageNet':
task_model = resnet.resnet18(num_classes=args.num_classes)
else:
print('WRONG DATASET!')
# loading pretrained
if args.pretrained:
print("Loading pretrained model", args.pretrained)
checkpoint = torch.load(args.pretrained)
task_model.load_state_dict({k: v for k, v in checkpoint['state_dict'].items() if 'fc' not in k}, strict=False) # copy all but last linear layers
#
vae = model.VAE(z_dim=args.latent_dim, nc=args.num_channels, s=args.arch_scaler)
discriminator = model.Discriminator(z_dim=args.latent_dim, s=args.arch_scaler)
#print("Sampling starts")
unlabeled_indices = np.setdiff1d(list(all_indices), current_indices)
unlabeled_sampler = data.sampler.SubsetRandomSampler(unlabeled_indices)
unlabeled_dataloader = data.DataLoader(train_dataset, sampler=unlabeled_sampler,
batch_size=args.batch_size, drop_last=False, **kwargs)
#print("Train starts")
# train the models on the current data
acc, vae, discriminator = solver.train(querry_dataloader,
task_model,
vae,
discriminator,
unlabeled_dataloader)
print('Final accuracy with {}% of data is: {:.2f}'.format(int(split*100.0*args.budget/args.num_images), acc))
accuracies.append(acc)
sampled_indices = solver.sample_for_labeling(vae, discriminator, unlabeled_dataloader)
current_indices = list(current_indices) + list(sampled_indices)
sampler = data.sampler.SubsetRandomSampler(current_indices)
querry_dataloader = data.DataLoader(train_dataset, sampler=sampler,
batch_size=args.batch_size, drop_last=False, **kwargs)
torch.save(accuracies, os.path.join(args.out_path, args.log_name))
def generate_model(opt):
assert opt.model in ['resnet']
if opt.model == 'resnet':
assert opt.model_depth in [10, 18, 34, 50, 101, 152, 200]
if opt.model_depth == 10:
model = resnet.resnet10(sample_input_W=opt.input_W,
sample_input_H=opt.input_H,
sample_input_D=opt.input_D,
shortcut_type=opt.resnet_shortcut,
no_cuda=opt.no_cuda,
num_seg_classes=opt.n_seg_classes)
elif opt.model_depth == 18:
model = resnet.resnet18(sample_input_W=opt.input_W,
sample_input_H=opt.input_H,
sample_input_D=opt.input_D,
shortcut_type=opt.resnet_shortcut,
no_cuda=opt.no_cuda,
num_seg_classes=opt.n_seg_classes)
elif opt.model_depth == 34:
model = resnet.resnet34(sample_input_W=opt.input_W,
sample_input_H=opt.input_H,
sample_input_D=opt.input_D,
shortcut_type=opt.resnet_shortcut,
no_cuda=opt.no_cuda,
num_seg_classes=opt.n_seg_classes)
elif opt.model_depth == 50:
model = resnet.resnet50(sample_input_W=opt.input_W,
sample_input_H=opt.input_H,
sample_input_D=opt.input_D,
shortcut_type=opt.resnet_shortcut,
no_cuda=opt.no_cuda,
num_seg_classes=opt.n_seg_classes)
elif opt.model_depth == 101:
model = resnet.resnet101(sample_input_W=opt.input_W,
sample_input_H=opt.input_H,
sample_input_D=opt.input_D,
shortcut_type=opt.resnet_shortcut,
no_cuda=opt.no_cuda,
num_seg_classes=opt.n_seg_classes)
elif opt.model_depth == 152:
model = resnet.resnet152(sample_input_W=opt.input_W,
sample_input_H=opt.input_H,
sample_input_D=opt.input_D,
shortcut_type=opt.resnet_shortcut,
no_cuda=opt.no_cuda,
num_seg_classes=opt.n_seg_classes)
elif opt.model_depth == 200:
model = resnet.resnet200(sample_input_W=opt.input_W,
sample_input_H=opt.input_H,
sample_input_D=opt.input_D,
shortcut_type=opt.resnet_shortcut,
no_cuda=opt.no_cuda,
num_seg_classes=opt.n_seg_classes)
if not opt.no_cuda:
if len(opt.gpu_id) > 1:
model = model.cuda()
model = nn.DataParallel(model, device_ids=opt.gpu_id)
net_dict = model.state_dict()
else:
import os
os.environ["CUDA_VISIBLE_DEVICES"] = str(opt.gpu_id[0])
model = model.cuda()
model = nn.DataParallel(model, device_ids=None)
net_dict = model.state_dict()
else:
net_dict = model.state_dict()
# load pretrain
if opt.phase != 'test' and opt.pretrain_path:
print('loading pretrained model {}'.format(opt.pretrain_path))
pretrain = torch.load(opt.pretrain_path)
pretrain_dict = {
k: v
for k, v in pretrain['state_dict'].items() if k in net_dict.keys()
}
net_dict.update(pretrain_dict)
model.load_state_dict(net_dict)
new_parameters = []
for pname, p in model.named_parameters():
for layer_name in opt.new_layer_names:
if pname.find(layer_name) >= 0:
new_parameters.append(p)
break
new_parameters_id = list(map(id, new_parameters))
base_parameters = list(
filter(lambda p: id(p) not in new_parameters_id,
model.parameters()))
parameters = {
'base_parameters': base_parameters,
'new_parameters': new_parameters
}
return model, parameters
return model, model.parameters()
def generate_model(opt):
assert opt.mode in ['score', 'feature']
if opt.mode == 'score':
last_fc = True
elif opt.mode == 'feature':
last_fc = False
assert opt.model_name in [
'resnet', 'preresnet', 'wideresnet', 'resnext', 'densenet'
]
if opt.model_name == 'resnet':
assert opt.model_depth in [10, 18, 34, 50, 101, 152, 200]
if opt.model_depth == 10:
model = resnet.resnet10(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_depth == 18:
model = resnet.resnet18(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_depth == 34:
model = resnet.resnet34(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_depth == 50:
model = resnet.resnet50(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_depth == 101:
model = resnet.resnet101(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_depth == 152:
model = resnet.resnet152(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_depth == 200:
model = resnet.resnet200(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_name == 'wideresnet':
assert opt.model_depth in [50]
if opt.model_depth == 50:
model = wide_resnet.resnet50(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
k=opt.wide_resnet_k,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_name == 'resnext':
assert opt.model_depth in [50, 101, 152]
if opt.model_depth == 50:
model = resnext.resnet50(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
cardinality=opt.resnext_cardinality,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_depth == 101:
model = resnext.resnet101(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
cardinality=opt.resnext_cardinality,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_depth == 152:
model = resnext.resnet152(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
cardinality=opt.resnext_cardinality,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_name == 'preresnet':
assert opt.model_depth in [18, 34, 50, 101, 152, 200]
if opt.model_depth == 18:
model = pre_act_resnet.resnet18(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_depth == 34:
model = pre_act_resnet.resnet34(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_depth == 50:
model = pre_act_resnet.resnet50(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_depth == 101:
model = pre_act_resnet.resnet101(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_depth == 152:
model = pre_act_resnet.resnet152(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_depth == 200:
model = pre_act_resnet.resnet200(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_name == 'densenet':
assert opt.model_depth in [121, 169, 201, 264]
if opt.model_depth == 121:
model = densenet.densenet121(num_classes=opt.n_classes,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_depth == 169:
model = densenet.densenet169(num_classes=opt.n_classes,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_depth == 201:
model = densenet.densenet201(num_classes=opt.n_classes,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_depth == 264:
model = densenet.densenet264(num_classes=opt.n_classes,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
if not opt.no_cuda:
model = model.cuda()
model = nn.DataParallel(model, device_ids=None)
return model
