## about model
num_classes = 20
## about data
data_dir = "../data/"
inupt_size = 224
batch_size = args.batch_size
## about training
num_epochs = args.num_epoch
lr = args.lr
gamma = args.gamma
## model initialization
model = models.model_C(num_classes=num_classes)
device = torch.device(args.cuda if torch.cuda.is_available() else "cpu")
model = model.to(device)
## data preparation
train_loader, valid_loader = data.load_data(data_dir=data_dir,
input_size=inupt_size,
batch_size=batch_size)
## optimizer
optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9)
## learning rate scheduler
milestones = [50 + i * 30 for i in range(10)]
scheduler = MultiStepLR(optimizer, milestones=milestones, gamma=gamma)
# scheduler = ExponentialLR(optimizer, gamma=gamma)
print("training: {:.4f}, {:.4f}".format(train_loss, train_acc))
print("train time: {:.4f} s".format(train_end - train_start))
if max_rate < train_acc:
max_rate = train_acc
if max_rate >= hyper.debug_threshold:
print("The model can overfit the training data, valid.")
else:
print("The model cannot overfit the training data, invalid.")
if __name__ == '__main__':
os.environ["CUDA_VISIBLE_DEVICES"] = "5"
## model initialization
model = models.model_C(num_classes=hyper.num_classes,
input_size=hyper.input_size)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = model.to(device)
## data preparation
_, _, debug_loader = data.load_data \
(data_dir = hyper.data_dir,input_size = hyper.input_size)
## optimizer
optimizer = optim.Adam(model.parameters(), lr=hyper.lr)
## loss function
criterion = nn.CrossEntropyLoss()
debug_model(model,