parser.add_argument('--data-dir', type=str, default="/ImageNet")
parser.add_argument('--batch-size', type=int, default=16)
parser.add_argument('--num-class', type=int, default=1000)
parser.add_argument('--num-epochs', type=int, default=100)
parser.add_argument('--lr', type=float, default=0.045)
parser.add_argument('--num-workers', type=int, default=0)
parser.add_argument('--gpus', type=str, default=0)
parser.add_argument('--print-freq', type=int, default=10)
parser.add_argument('--save-epoch-freq', type=int, default=1)
parser.add_argument('--save-path', type=str, default="output")
parser.add_argument('--resume', type=str, default="", help="For training from one checkpoint")
parser.add_argument('--start-epoch', type=int, default=0, help="Corresponding to the epoch of resume ")
args = parser.parse_args()
# read data
dataloders, dataset_sizes = ImageNetData(args)
# use gpu or not
use_gpu = torch.cuda.is_available()
print("use_gpu:{}".format(use_gpu))
# get model
model = mobilenetv2_19(num_classes = args.num_class)
if args.resume:
if os.path.isfile(args.resume):
print(("=> loading checkpoint '{}'".format(args.resume)))
checkpoint = torch.load(args.resume)
base_dict = {'.'.join(k.split('.')[1:]): v for k, v in list(checkpoint.state_dict().items())}
model.load_state_dict(base_dict)
else:
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])
# train_data = datasets.ImageFolder('data/dogs-vs-cats/train', train_transforms)
# valid_data = datasets.ImageFolder('data/dogs-vs-cats/valid', test_transforms)
# test_data = datasets.ImageFolder('data/dogs-vs-cats/test', test_transforms)
BATCH_SIZE = 16
# train_iterator = torch.utils.data.DataLoader(train_data, shuffle=True, batch_size=BATCH_SIZE)
# valid_iterator = torch.utils.data.DataLoader(valid_data, batch_size=BATCH_SIZE)
# test_iterator = torch.utils.data.DataLoader(test_data, batch_size=BATCH_SIZE)
dataloders, dataset_sizes = ImageNetData(data_dir="ImageData",
batch_size=BATCH_SIZE,
num_workers=0)
train_iterator = dataloders['train']
valid_iterator = dataloders['val']
device = torch.device('cuda')
model = models.resnet50(pretrained=True).to(device)
for param in model.parameters():
param.requires_grad = False
model.fc = nn.Linear(in_features=2048, out_features=54).to(device)
optimizer = optim.Adam(model.parameters())
criterion = nn.CrossEntropyLoss()
EPOCHS = 30