| 12 | return tree |
| 13 | |
| 14 | def init_tree(tree: ProtoTree, optimizer, scheduler, device, args: argparse.Namespace): |
| 15 | epoch = 1 |
| 16 | mean = 0.5 |
| 17 | std = 0.1 |
| 18 | # load trained prototree if flag is set |
| 19 | |
| 20 | # NOTE: TRAINING FURTHER FROM A CHECKPOINT DOESN'T SEEM TO WORK CORRECTLY. EVALUATING A TRAINED PROTOTREE FROM A CHECKPOINT DOES WORK. |
| 21 | if args.state_dict_dir_tree != '': |
| 22 | if not args.disable_cuda and torch.cuda.is_available(): |
| 23 | device = torch.device('cuda:{}'.format(torch.cuda.current_device())) |
| 24 | else: |
| 25 | device = torch.device('cpu') |
| 26 | |
| 27 | |
| 28 | if args.disable_cuda or not torch.cuda.is_available(): |
| 29 | # tree = load_state(args.state_dict_dir_tree, device) |
| 30 | tree = torch.load(args.state_dict_dir_tree+'/model.pth', map_location=device) |
| 31 | else: |
| 32 | tree = torch.load(args.state_dict_dir_tree+'/model.pth') |
| 33 | tree.to(device=device) |
| 34 | try: |
| 35 | epoch = int(args.state_dict_dir_tree.split('epoch_')[-1]) + 1 |
| 36 | except: |
| 37 | epoch=args.epochs+1 |
| 38 | print("Train further from epoch: ", epoch, flush=True) |
| 39 | optimizer.load_state_dict(torch.load(args.state_dict_dir_tree+'/optimizer_state.pth', map_location=device)) |
| 40 | |
| 41 | if epoch>args.freeze_epochs: |
| 42 | for parameter in tree._net.parameters(): |
| 43 | parameter.requires_grad = True |
| 44 | if not args.disable_derivative_free_leaf_optim: |
| 45 | for leaf in tree.leaves: |
| 46 | leaf._dist_params.requires_grad = False |
| 47 | |
| 48 | if os.path.isfile(args.state_dict_dir_tree+'/scheduler_state.pth'): |
| 49 | # scheduler.load_state_dict(torch.load(args.state_dict_dir_tree+'/scheduler_state.pth')) |
| 50 | # print(scheduler.state_dict(),flush=True) |
| 51 | scheduler.last_epoch = epoch - 1 |
| 52 | scheduler._step_count = epoch |
| 53 | |
| 54 | |
| 55 | elif args.state_dict_dir_net != '': # load pretrained conv network |
| 56 | # initialize prototypes |
| 57 | torch.nn.init.normal_(tree.prototype_layer.prototype_vectors, mean=mean, std=std) |
| 58 | #strict is False so when loading pretrained model, ignore the linear classification layer |
| 59 | tree._net.load_state_dict(torch.load(args.state_dict_dir_net+'/model_state.pth'), strict=False) |
| 60 | tree._add_on.load_state_dict(torch.load(args.state_dict_dir_net+'/model_state.pth'), strict=False) |
| 61 | else: |
| 62 | with torch.no_grad(): |
| 63 | # initialize prototypes |
| 64 | torch.nn.init.normal_(tree.prototype_layer.prototype_vectors, mean=mean, std=std) |
| 65 | tree._add_on.apply(init_weights_xavier) |
| 66 | return tree, epoch |
| 67 | |
| 68 | def init_weights_xavier(m): |
| 69 | if type(m) == torch.nn.Conv2d: |