lichen/ieemoo-ai-conpurchase
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

udpate

Browse Source
This commit is contained in:
李晨
2023-08-17 11:21:39 +08:00
parent 396cc892d5
commit a56f93d2b5
2 changed files with 1 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
ieemoo-ai-conpurchase.py
View File

@ -10,7 +10,7 @@ from utils.embedding import DataProcessing as dp
app = Flask(__name__) app = Flask(__name__)
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument('--model', default='./checkpoint/raft-small.pth',help="restore checkpoint") parser.add_argument('--model', default='../module/ieemoo-ai-conpurchase/model/now/raft-small.pth',help="restore checkpoint")
parser.add_argument('--checkpoint', default='mobilevit',help="get embedding ") parser.add_argument('--checkpoint', default='mobilevit',help="get embedding ")
parser.add_argument('--device', default='cuda',help="device") parser.add_argument('--device', default='cuda',help="device")
parser.add_argument('--small', type=bool, default=True, help='use small model') parser.add_argument('--small', type=bool, default=True, help='use small model')

1
network/BaseNet.py
View File

@ -170,7 +170,6 @@ if __name__ == '__main__':
img = transform(img) img = transform(img)
img = img.cuda() img = img.cuda()
# from torchsummary import summary # from torchsummary import summary
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = ResnetFpn().to(device) model = ResnetFpn().to(device)